refactor

.ci/update_windows/update.py

@@ -53,16 +53,6 @@ try:
     repo.stash(ident)
 except KeyError:
     print("nothing to stash")  # noqa: T201
-except:
-    print("Could not stash, cleaning index and trying again.")  # noqa: T201
-    repo.state_cleanup()
-    repo.index.read_tree(repo.head.peel().tree)
-    repo.index.write()
-    try:
-        repo.stash(ident)
-    except KeyError:
-        print("nothing to stash.")  # noqa: T201
-
 backup_branch_name = 'backup_branch_{}'.format(datetime.today().strftime('%Y-%m-%d_%H_%M_%S'))
 print("creating backup branch: {}".format(backup_branch_name))  # noqa: T201
 try:
@@ -76,10 +66,8 @@ if branch is None:
     try:
         ref = repo.lookup_reference('refs/remotes/origin/master')
     except:
-        print("fetching.")  # noqa: T201
-        for remote in repo.remotes:
-            if remote.name == "origin":
-                remote.fetch()
+        print("pulling.")  # noqa: T201
+        pull(repo)
         ref = repo.lookup_reference('refs/remotes/origin/master')
     repo.checkout(ref)
     branch = repo.lookup_branch('master')
@@ -161,4 +149,3 @@ try:
         shutil.copy(stable_update_script, stable_update_script_to)
 except:
     pass
-

.ci/windows_amd_base_files/README_VERY_IMPORTANT.txt

@@ -1,5 +1,5 @@
-As of the time of writing this you need this driver for best results:
-https://www.amd.com/en/resources/support-articles/release-notes/RN-AMDGPU-WINDOWS-PYTORCH-7-1-1.html
+As of the time of writing this you need this preview driver for best results:
+https://www.amd.com/en/resources/support-articles/release-notes/RN-AMDGPU-WINDOWS-PYTORCH-PREVIEW.html
 
 HOW TO RUN:
 
@@ -25,4 +25,3 @@ In the ComfyUI directory you will find a file: extra_model_paths.yaml.example
 Rename this file to: extra_model_paths.yaml and edit it with your favorite text editor.
 
 
-

.github/PULL_REQUEST_TEMPLATE/api-node.md

@@ -1,21 +0,0 @@
-<!-- API_NODE_PR_CHECKLIST: do not remove -->
-
-## API Node PR Checklist
-
-### Scope
-- [ ] **Is API Node Change**
-
-### Pricing & Billing
-- [ ] **Need pricing update**
-- [ ] **No pricing update**
-
-If **Need pricing update**:
-- [ ] Metronome rate cards updated
-- [ ] Auto‑billing tests updated and passing
-
-### QA
-- [ ] **QA done**
-- [ ] **QA not required**
-
-### Comms
-- [ ] Informed **Kosinkadink**

.github/workflows/api-node-template.yml

@@ -1,58 +0,0 @@
-name: Append API Node PR template
-
-on:
-  pull_request_target:
-    types: [opened, reopened, synchronize, ready_for_review]
-    paths:
-      - 'comfy_api_nodes/**'   # only run if these files changed
-
-permissions:
-  contents: read
-  pull-requests: write
-
-jobs:
-  inject:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Ensure template exists and append to PR body
-        uses: actions/github-script@v7
-        with:
-          script: |
-            const { owner, repo } = context.repo;
-            const number = context.payload.pull_request.number;
-            const templatePath = '.github/PULL_REQUEST_TEMPLATE/api-node.md';
-            const marker = '<!-- API_NODE_PR_CHECKLIST: do not remove -->';
-
-            const { data: pr } = await github.rest.pulls.get({ owner, repo, pull_number: number });
-
-            let templateText;
-            try {
-              const res = await github.rest.repos.getContent({
-                owner,
-                repo,
-                path: templatePath,
-                ref: pr.base.ref
-              });
-              const buf = Buffer.from(res.data.content, res.data.encoding || 'base64');
-              templateText = buf.toString('utf8');
-            } catch (e) {
-              core.setFailed(`Required PR template not found at "${templatePath}" on ${pr.base.ref}. Please add it to the repo.`);
-              return;
-            }
-
-            // Enforce the presence of the marker inside the template (for idempotence)
-            if (!templateText.includes(marker)) {
-              core.setFailed(`Template at "${templatePath}" does not contain the required marker:\n${marker}\nAdd it so we can detect duplicates safely.`);
-              return;
-            }
-
-            // If the PR already contains the marker, do not append again.
-            const body = pr.body || '';
-            if (body.includes(marker)) {
-              core.info('Template already present in PR body; nothing to inject.');
-              return;
-            }
-
-            const newBody = (body ? body + '\n\n' : '') + templateText + '\n';
-            await github.rest.pulls.update({ owner, repo, pull_number: number, body: newBody });
-            core.notice('API Node template appended to PR description.');

.github/workflows/release-stable-all.yml

@@ -14,7 +14,7 @@ jobs:
       contents: "write"
       packages: "write"
       pull-requests: "read"
-    name: "Release NVIDIA Default (cu130)"
+    name: "Release NVIDIA Default (cu129)"
     uses: ./.github/workflows/stable-release.yml
     with:
       git_tag: ${{ inputs.git_tag }}
@@ -43,33 +43,16 @@ jobs:
       test_release: true
     secrets: inherit
 
-  release_nvidia_cu126:
-    permissions:
-      contents: "write"
-      packages: "write"
-      pull-requests: "read"
-    name: "Release NVIDIA cu126"
-    uses: ./.github/workflows/stable-release.yml
-    with:
-      git_tag: ${{ inputs.git_tag }}
-      cache_tag: "cu126"
-      python_minor: "12"
-      python_patch: "10"
-      rel_name: "nvidia"
-      rel_extra_name: "_cu126"
-      test_release: true
-    secrets: inherit
-
   release_amd_rocm:
     permissions:
       contents: "write"
       packages: "write"
       pull-requests: "read"
-    name: "Release AMD ROCm 7.1.1"
+    name: "Release AMD ROCm 6.4.4"
     uses: ./.github/workflows/stable-release.yml
     with:
       git_tag: ${{ inputs.git_tag }}
-      cache_tag: "rocm711"
+      cache_tag: "rocm644"
       python_minor: "12"
       python_patch: "10"
       rel_name: "amd"

.github/workflows/test-ci.yml

@@ -5,7 +5,6 @@ on:
   push:
     branches:
       - master
-      - release/**
     paths-ignore:
       - 'app/**'
       - 'input/**'
@@ -22,15 +21,14 @@ jobs:
       fail-fast: false
       matrix:
         # os: [macos, linux, windows]
-        # os: [macos, linux]
-        os: [linux]
-        python_version: ["3.10", "3.11", "3.12"]
+        os: [macos, linux]
+        python_version: ["3.9", "3.10", "3.11", "3.12"]
         cuda_version: ["12.1"]
         torch_version: ["stable"]
         include:
-          # - os: macos
-          #   runner_label: [self-hosted, macOS]
-          #   flags: "--use-pytorch-cross-attention"
+          - os: macos
+            runner_label: [self-hosted, macOS]
+            flags: "--use-pytorch-cross-attention"
           - os: linux
             runner_label: [self-hosted, Linux]
             flags: ""
@@ -75,15 +73,14 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        # os: [macos, linux]
-        os: [linux]
+        os: [macos, linux]
         python_version: ["3.11"]
         cuda_version: ["12.1"]
         torch_version: ["nightly"]
         include:
-          # - os: macos
-          #   runner_label: [self-hosted, macOS]
-          #   flags: "--use-pytorch-cross-attention"
+          - os: macos
+            runner_label: [self-hosted, macOS]
+            flags: "--use-pytorch-cross-attention"
           - os: linux
             runner_label: [self-hosted, Linux]
             flags: ""

.github/workflows/test-execution.yml

@@ -2,9 +2,9 @@ name: Execution Tests
 
 on:
   push:
-    branches: [ main, master, release/** ]
+    branches: [ main, master ]
   pull_request:
-    branches: [ main, master, release/** ]
+    branches: [ main, master ]
 
 jobs:
   test:

.github/workflows/test-launch.yml

@@ -2,9 +2,9 @@ name: Test server launches without errors
 
 on:
   push:
-    branches: [ main, master, release/** ]
+    branches: [ main, master ]
   pull_request:
-    branches: [ main, master, release/** ]
+    branches: [ main, master ]
 
 jobs:
   test:

.github/workflows/test-unit.yml

@@ -2,9 +2,9 @@ name: Unit Tests
 
 on:
   push:
-    branches: [ main, master, release/** ]
+    branches: [ main, master ]
   pull_request:
-    branches: [ main, master, release/** ]
+    branches: [ main, master ]
 
 jobs:
   test:

.github/workflows/update-version.yml

@@ -6,7 +6,6 @@ on:
       - "pyproject.toml"
     branches:
       - master
-      - release/**
 
 jobs:
   update-version:

CODEOWNERS

@@ -1,2 +1,3 @@
 # Admins
-* @comfyanonymous @kosinkadink @guill
+* @comfyanonymous
+* @kosinkadink

QUANTIZATION.md

@@ -1,168 +0,0 @@
-# The Comfy guide to Quantization
-
-
-## How does quantization work?
-
-Quantization aims to map a high-precision value x_f to a lower precision format with minimal loss in accuracy. These smaller formats then serve to reduce the models memory footprint and increase throughput by using specialized hardware.
-
-When simply converting a value from FP16 to FP8 using the round-nearest method we might hit two issues:
-- The dynamic range of FP16 (-65,504, 65,504) far exceeds FP8 formats like E4M3 (-448, 448) or E5M2 (-57,344, 57,344), potentially resulting in clipped values
-- The original values are concentrated in a small range (e.g. -1,1) leaving many FP8-bits "unused"
-
-By using a scaling factor, we aim to map these values into the quantized-dtype range, making use of the full spectrum. One of the easiest approaches, and common, is using per-tensor absolute-maximum scaling.
-
-```
-absmax = max(abs(tensor))
-scale = amax / max_dynamic_range_low_precision
-
-# Quantization
-tensor_q = (tensor / scale).to(low_precision_dtype)
-
-# De-Quantization
-tensor_dq = tensor_q.to(fp16) * scale
-
-tensor_dq ~ tensor
-```
-
-Given that additional information (scaling factor) is needed to "interpret" the quantized values, we describe those as derived datatypes.
-
-
-## Quantization in Comfy
-
-```
-QuantizedTensor (torch.Tensor subclass)
-  ↓ __torch_dispatch__
-Two-Level Registry (generic + layout handlers)
-  ↓
-MixedPrecisionOps + Metadata Detection
-```
-
-### Representation
-
-To represent these derived datatypes, ComfyUI uses a subclass of torch.Tensor to implements these using the `QuantizedTensor` class found in `comfy/quant_ops.py`
-
-A `Layout` class defines how a specific quantization format behaves:
-- Required parameters
-- Quantize method
-- De-Quantize method
-
-```python
-from comfy.quant_ops import QuantizedLayout
-
-class MyLayout(QuantizedLayout):
-    @classmethod
-    def quantize(cls, tensor, **kwargs):
-        # Convert to quantized format
-        qdata = ...
-        params = {'scale': ..., 'orig_dtype': tensor.dtype}
-        return qdata, params
-    
-    @staticmethod
-    def dequantize(qdata, scale, orig_dtype, **kwargs):
-        return qdata.to(orig_dtype) * scale
-```
-
-To then run operations using these QuantizedTensors we use two registry systems to define supported operations. 
-The first is a **generic registry** that handles operations common to all quantized formats (e.g., `.to()`, `.clone()`, `.reshape()`).
-
-The second registry is layout-specific and allows to implement fast-paths like nn.Linear.
-```python
-from comfy.quant_ops import register_layout_op
-
-@register_layout_op(torch.ops.aten.linear.default, MyLayout)
-def my_linear(func, args, kwargs):
-    # Extract tensors, call optimized kernel
-    ...
-```
-When `torch.nn.functional.linear()` is called with QuantizedTensor arguments, `__torch_dispatch__` automatically routes to the registered implementation.
-For any unsupported operation, QuantizedTensor will fallback to call `dequantize` and dispatch using the high-precision implementation.
-
-
-### Mixed Precision
-
-The `MixedPrecisionOps` class (lines 542-648 in `comfy/ops.py`) enables per-layer quantization decisions, allowing different layers in a model to use different precisions. This is activated when a model config contains a `layer_quant_config` dictionary that specifies which layers should be quantized and how.
-
-**Architecture:**
-
-```python
-class MixedPrecisionOps(disable_weight_init):
-    _layer_quant_config = {}  # Maps layer names to quantization configs
-    _compute_dtype = torch.bfloat16  # Default compute / dequantize precision
-```
-
-**Key mechanism:**
-
-The custom `Linear._load_from_state_dict()` method inspects each layer during model loading:
-- If the layer name is **not** in `_layer_quant_config`: load weight as regular tensor in `_compute_dtype`
-- If the layer name **is** in `_layer_quant_config`: 
-  - Load weight as `QuantizedTensor` with the specified layout (e.g., `TensorCoreFP8Layout`)
-  - Load associated quantization parameters (scales, block_size, etc.)
-
-**Why it's needed:**
-
-Not all layers tolerate quantization equally. Sensitive operations like final projections can be kept in higher precision, while compute-heavy matmuls are quantized. This provides most of the performance benefits while maintaining quality.
-
-The system is selected in `pick_operations()` when `model_config.layer_quant_config` is present, making it the highest-priority operation mode.
-
-
-## Checkpoint Format
-
-Quantized checkpoints are stored as standard safetensors files with quantized weight tensors and associated scaling parameters, plus a `_quantization_metadata` JSON entry describing the quantization scheme.
-
-The quantized checkpoint will contain the same layers as the original checkpoint but:
-- The weights are stored as quantized values, sometimes using a different storage datatype. E.g. uint8 container for fp8.
-- For each quantized weight a number of additional scaling parameters are stored alongside depending on the recipe.
-- We store a metadata.json in the metadata of the final safetensor containing the `_quantization_metadata` describing which layers are quantized and what layout has been used.
-
-### Scaling Parameters details
-We define 4 possible scaling parameters that should cover most recipes in the near-future:
-- **weight_scale**: quantization scalers for the weights
-- **weight_scale_2**: global scalers in the context of double scaling
-- **pre_quant_scale**: scalers used for smoothing salient weights
-- **input_scale**: quantization scalers for the activations
-
-| Format | Storage dtype | weight_scale | weight_scale_2 | pre_quant_scale | input_scale |
-|--------|---------------|--------------|----------------|-----------------|-------------|
-| float8_e4m3fn | float32 | float32 (scalar) | - | - | float32 (scalar) |
-
-You can find the defined formats in `comfy/quant_ops.py` (QUANT_ALGOS).
-
-### Quantization Metadata
-
-The metadata stored alongside the checkpoint contains:
-- **format_version**: String to define a version of the standard
-- **layers**: A dictionary mapping layer names to their quantization format. The format string maps to the definitions found in `QUANT_ALGOS`. 
-
-Example:
-```json
-{
-  "_quantization_metadata": {
-    "format_version": "1.0",
-    "layers": {
-      "model.layers.0.mlp.up_proj": "float8_e4m3fn",
-      "model.layers.0.mlp.down_proj": "float8_e4m3fn",
-      "model.layers.1.mlp.up_proj": "float8_e4m3fn"
-    }
-  }
-}
-```
-
-
-## Creating Quantized Checkpoints
-
-To create compatible checkpoints, use any quantization tool provided the output follows the checkpoint format described above and uses a layout defined in `QUANT_ALGOS`.
-
-### Weight Quantization
-
-Weight quantization is straightforward - compute the scaling factor directly from the weight tensor using the absolute maximum method described earlier. Each layer's weights are quantized independently and stored with their corresponding `weight_scale` parameter.
-
-### Calibration (for Activation Quantization)
-
-Activation quantization (e.g., for FP8 Tensor Core operations) requires `input_scale` parameters that cannot be determined from static weights alone. Since activation values depend on actual inputs, we use **post-training calibration (PTQ)**:
-
-1. **Collect statistics**: Run inference on N representative samples
-2. **Track activations**: Record the absolute maximum (`amax`) of inputs to each quantized layer
-3. **Compute scales**: Derive `input_scale` from collected statistics
-4. **Store in checkpoint**: Save `input_scale` parameters alongside weights
-
-The calibration dataset should be representative of your target use case. For diffusion models, this typically means a diverse set of prompts and generation parameters.

README.md

@@ -67,8 +67,6 @@ See what ComfyUI can do with the [example workflows](https://comfyanonymous.gith
    - [HiDream](https://comfyanonymous.github.io/ComfyUI_examples/hidream/)
    - [Qwen Image](https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/)
    - [Hunyuan Image 2.1](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_image/)
-   - [Flux 2](https://comfyanonymous.github.io/ComfyUI_examples/flux2/)
-   - [Z Image](https://comfyanonymous.github.io/ComfyUI_examples/z_image/)
 - Image Editing Models
    - [Omnigen 2](https://comfyanonymous.github.io/ComfyUI_examples/omnigen/)
    - [Flux Kontext](https://comfyanonymous.github.io/ComfyUI_examples/flux/#flux-kontext-image-editing-model)
@@ -81,7 +79,6 @@ See what ComfyUI can do with the [example workflows](https://comfyanonymous.gith
    - [Hunyuan Video](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_video/)
    - [Wan 2.1](https://comfyanonymous.github.io/ComfyUI_examples/wan/)
    - [Wan 2.2](https://comfyanonymous.github.io/ComfyUI_examples/wan22/)
-   - [Hunyuan Video 1.5](https://docs.comfy.org/tutorials/video/hunyuan/hunyuan-video-1-5)
 - Audio Models
    - [Stable Audio](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
    - [ACE Step](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
@@ -119,9 +116,6 @@ ComfyUI follows a weekly release cycle targeting Monday but this regularly chang
 
 1. **[ComfyUI Core](https://github.com/comfyanonymous/ComfyUI)**
    - Releases a new stable version (e.g., v0.7.0) roughly every week.
-   - Starting from v0.4.0 patch versions will be used for fixes backported onto the current stable release.
-   - Minor versions will be used for releases off the master branch.
-   - Patch versions may still be used for releases on the master branch in cases where a backport would not make sense.
    - Commits outside of the stable release tags may be very unstable and break many custom nodes.
    - Serves as the foundation for the desktop release
 
@@ -179,7 +173,7 @@ There is a portable standalone build for Windows that should work for running on
 
 ### [Direct link to download](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia.7z)
 
-Simply download, extract with [7-Zip](https://7-zip.org) or with the windows explorer on recent windows versions and run. For smaller models you normally only need to put the checkpoints (the huge ckpt/safetensors files) in: ComfyUI\models\checkpoints but many of the larger models have multiple files. Make sure to follow the instructions to know which subfolder to put them in ComfyUI\models\
+Simply download, extract with [7-Zip](https://7-zip.org) and run. Make sure you put your Stable Diffusion checkpoints/models (the huge ckpt/safetensors files) in: ComfyUI\models\checkpoints
 
 If you have trouble extracting it, right click the file -> properties -> unblock
 
@@ -189,9 +183,7 @@ Update your Nvidia drivers if it doesn't start.
 
 [Experimental portable for AMD GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_amd.7z)
 
-[Portable with pytorch cuda 12.8 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu128.7z).
-
-[Portable with pytorch cuda 12.6 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu126.7z) (Supports Nvidia 10 series and older GPUs).
+[Portable with pytorch cuda 12.8 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu128.7z) (Supports Nvidia 10 series and older GPUs).
 
 #### How do I share models between another UI and ComfyUI?
 
@@ -208,7 +200,7 @@ comfy install
 
 ## Manual Install (Windows, Linux)
 
-Python 3.14 works but you may encounter issues with the torch compile node. The free threaded variant is still missing some dependencies.
+Python 3.14 will work if you comment out the `kornia` dependency in the requirements.txt file (breaks the canny node) but it is not recommended.
 
 Python 3.13 is very well supported. If you have trouble with some custom node dependencies on 3.13 you can try 3.12
 
@@ -229,7 +221,7 @@ AMD users can install rocm and pytorch with pip if you don't have it already ins
 
 This is the command to install the nightly with ROCm 7.0 which might have some performance improvements:
 
-```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm7.1```
+```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm7.0```
 
 
 ### AMD GPUs (Experimental: Windows and Linux), RDNA 3, 3.5 and 4 only.
@@ -250,7 +242,7 @@ RDNA 4 (RX 9000 series):
 
 ### Intel GPUs (Windows and Linux)
 
-Intel Arc GPU users can install native PyTorch with torch.xpu support using pip. More information can be found [here](https://pytorch.org/docs/main/notes/get_start_xpu.html)
+(Option 1) Intel Arc GPU users can install native PyTorch with torch.xpu support using pip. More information can be found [here](https://pytorch.org/docs/main/notes/get_start_xpu.html)
 
 1. To install PyTorch xpu, use the following command:
 
@@ -260,6 +252,10 @@ This is the command to install the Pytorch xpu nightly which might have some per
 
 ```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/xpu```
 
+(Option 2) Alternatively, Intel GPUs supported by Intel Extension for PyTorch (IPEX) can leverage IPEX for improved performance.
+
+1. visit [Installation](https://intel.github.io/intel-extension-for-pytorch/index.html#installation?platform=gpu) for more information.
+
 ### NVIDIA
 
 Nvidia users should install stable pytorch using this command:
@@ -323,32 +319,6 @@ For models compatible with Iluvatar Extension for PyTorch. Here's a step-by-step
 1. Install the Iluvatar Corex Toolkit by adhering to the platform-specific instructions on the [Installation](https://support.iluvatar.com/#/DocumentCentre?id=1&nameCenter=2&productId=520117912052801536)
 2. Launch ComfyUI by running `python main.py`
 
-
-## [ComfyUI-Manager](https://github.com/Comfy-Org/ComfyUI-Manager/tree/manager-v4)
-
-**ComfyUI-Manager** is an extension that allows you to easily install, update, and manage custom nodes for ComfyUI.
-
-### Setup
-
-1. Install the manager dependencies:
-   ```bash
-   pip install -r manager_requirements.txt
-   ```
-
-2. Enable the manager with the `--enable-manager` flag when running ComfyUI:
-   ```bash
-   python main.py --enable-manager
-   ```
-
-### Command Line Options
-
-| Flag | Description |
-|------|-------------|
-| `--enable-manager` | Enable ComfyUI-Manager |
-| `--enable-manager-legacy-ui` | Use the legacy manager UI instead of the new UI (requires `--enable-manager`) |
-| `--disable-manager-ui` | Disable the manager UI and endpoints while keeping background features like security checks and scheduled installation completion (requires `--enable-manager`) |
-
-
 # Running
 
 ```python main.py```

api_server/routes/internal/internal_routes.py

@@ -58,13 +58,8 @@ class InternalRoutes:
                 return web.json_response({"error": "Invalid directory type"}, status=400)
 
             directory = get_directory_by_type(directory_type)
-
-            def is_visible_file(entry: os.DirEntry) -> bool:
-                """Filter out hidden files (e.g., .DS_Store on macOS)."""
-                return entry.is_file() and not entry.name.startswith('.')
-
             sorted_files = sorted(
-                (entry for entry in os.scandir(directory) if is_visible_file(entry)),
+                (entry for entry in os.scandir(directory) if entry.is_file()),
                 key=lambda entry: -entry.stat().st_mtime
             )
             return web.json_response([entry.name for entry in sorted_files], status=200)

app/user_manager.py

@@ -59,9 +59,6 @@ class UserManager():
         user = "default"
         if args.multi_user and "comfy-user" in request.headers:
             user = request.headers["comfy-user"]
-            # Block System Users (use same error message to prevent probing)
-            if user.startswith(folder_paths.SYSTEM_USER_PREFIX):
-                raise KeyError("Unknown user: " + user)
 
         if user not in self.users:
             raise KeyError("Unknown user: " + user)
@@ -69,16 +66,15 @@ class UserManager():
         return user
 
     def get_request_user_filepath(self, request, file, type="userdata", create_dir=True):
+        user_directory = folder_paths.get_user_directory()
+
         if type == "userdata":
-            root_dir = folder_paths.get_user_directory()
+            root_dir = user_directory
         else:
             raise KeyError("Unknown filepath type:" + type)
 
         user = self.get_request_user_id(request)
-        user_root = folder_paths.get_public_user_directory(user)
-        if user_root is None:
-            return None
-        path = user_root
+        path = user_root = os.path.abspath(os.path.join(root_dir, user))
 
         # prevent leaving /{type}
         if os.path.commonpath((root_dir, user_root)) != root_dir:
@@ -105,11 +101,7 @@ class UserManager():
         name = name.strip()
         if not name:
             raise ValueError("username not provided")
-        if name.startswith(folder_paths.SYSTEM_USER_PREFIX):
-            raise ValueError("System User prefix not allowed")
         user_id = re.sub("[^a-zA-Z0-9-_]+", '-', name)
-        if user_id.startswith(folder_paths.SYSTEM_USER_PREFIX):
-            raise ValueError("System User prefix not allowed")
         user_id = user_id + "_" + str(uuid.uuid4())
 
         self.users[user_id] = name
@@ -140,10 +132,7 @@ class UserManager():
             if username in self.users.values():
                 return web.json_response({"error": "Duplicate username."}, status=400)
 
-            try:
-                user_id = self.add_user(username)
-            except ValueError as e:
-                return web.json_response({"error": str(e)}, status=400)
+            user_id = self.add_user(username)
             return web.json_response(user_id)
 
         @routes.get("/userdata")
@@ -435,7 +424,7 @@ class UserManager():
                 return source
 
             dest = get_user_data_path(request, check_exists=False, param="dest")
-            if not isinstance(dest, str):
+            if not isinstance(source, str):
                 return dest
 
             overwrite = request.query.get("overwrite", 'true') != "false"

comfy/cldm/cldm.py

@@ -413,8 +413,7 @@ class ControlNet(nn.Module):
         out_middle = []
 
         if self.num_classes is not None:
-            if y is None:
-                raise ValueError("y is None, did you try using a controlnet for SDXL on SD1?")
+            assert y.shape[0] == x.shape[0]
             emb = emb + self.label_emb(y)
 
         h = x

comfy/cli_args.py

@@ -97,13 +97,6 @@ class LatentPreviewMethod(enum.Enum):
     Latent2RGB = "latent2rgb"
     TAESD = "taesd"
 
-    @classmethod
-    def from_string(cls, value: str):
-        for member in cls:
-            if member.value == value:
-                return member
-        return None
-
 parser.add_argument("--preview-method", type=LatentPreviewMethod, default=LatentPreviewMethod.NoPreviews, help="Default preview method for sampler nodes.", action=EnumAction)
 
 parser.add_argument("--preview-size", type=int, default=512, help="Sets the maximum preview size for sampler nodes.")
@@ -128,12 +121,6 @@ upcast.add_argument("--force-upcast-attention", action="store_true", help="Force
 upcast.add_argument("--dont-upcast-attention", action="store_true", help="Disable all upcasting of attention. Should be unnecessary except for debugging.")
 
 
-parser.add_argument("--enable-manager", action="store_true", help="Enable the ComfyUI-Manager feature.")
-manager_group = parser.add_mutually_exclusive_group()
-manager_group.add_argument("--disable-manager-ui", action="store_true", help="Disables only the ComfyUI-Manager UI and endpoints. Scheduled installations and similar background tasks will still operate.")
-manager_group.add_argument("--enable-manager-legacy-ui", action="store_true", help="Enables the legacy UI of ComfyUI-Manager")
-
-
 vram_group = parser.add_mutually_exclusive_group()
 vram_group.add_argument("--gpu-only", action="store_true", help="Store and run everything (text encoders/CLIP models, etc... on the GPU).")
 vram_group.add_argument("--highvram", action="store_true", help="By default models will be unloaded to CPU memory after being used. This option keeps them in GPU memory.")
@@ -144,8 +131,7 @@ vram_group.add_argument("--cpu", action="store_true", help="To use the CPU for e
 
 parser.add_argument("--reserve-vram", type=float, default=None, help="Set the amount of vram in GB you want to reserve for use by your OS/other software. By default some amount is reserved depending on your OS.")
 
-parser.add_argument("--async-offload", nargs='?', const=2, type=int, default=None, metavar="NUM_STREAMS", help="Use async weight offloading. An optional argument controls the amount of offload streams. Default is 2. Enabled by default on Nvidia.")
-parser.add_argument("--disable-async-offload", action="store_true", help="Disable async weight offloading.")
+parser.add_argument("--async-offload", action="store_true", help="Use async weight offloading.")
 
 parser.add_argument("--force-non-blocking", action="store_true", help="Force ComfyUI to use non-blocking operations for all applicable tensors. This may improve performance on some non-Nvidia systems but can cause issues with some workflows.")
 
@@ -174,14 +160,13 @@ parser.add_argument("--windows-standalone-build", action="store_true", help="Win
 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("--whitelist-custom-nodes", type=str, nargs='+', default=[], help="Specify custom node folders to load even when --disable-all-custom-nodes is enabled.")
-parser.add_argument("--disable-api-nodes", action="store_true", help="Disable loading all api nodes. Also prevents the frontend from communicating with the internet.")
+parser.add_argument("--disable-api-nodes", action="store_true", help="Disable loading all api nodes.")
 
 parser.add_argument("--multi-user", action="store_true", help="Enables per-user storage.")
 
 parser.add_argument("--verbose", default='INFO', const='DEBUG', nargs="?", choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], help='Set the logging level')
 parser.add_argument("--log-stdout", action="store_true", help="Send normal process output to stdout instead of stderr (default).")
 
-
 # The default built-in provider hosted under web/
 DEFAULT_VERSION_STRING = "comfyanonymous/ComfyUI@latest"
 

comfy/context_windows.py

@@ -51,43 +51,32 @@ class ContextHandlerABC(ABC):
 
 
 class IndexListContextWindow(ContextWindowABC):
-    def __init__(self, index_list: list[int], dim: int=0, total_frames: int=0):
+    def __init__(self, index_list: list[int], dim: int=0):
         self.index_list = index_list
         self.context_length = len(index_list)
         self.dim = dim
-        self.total_frames = total_frames
-        self.center_ratio = (min(index_list) + max(index_list)) / (2 * total_frames)
 
-    def get_tensor(self, full: torch.Tensor, device=None, dim=None, retain_index_list=[]) -> torch.Tensor:
+    def get_tensor(self, full: torch.Tensor, device=None, dim=None) -> torch.Tensor:
         if dim is None:
             dim = self.dim
         if dim == 0 and full.shape[dim] == 1:
             return full
-        idx = tuple([slice(None)] * dim + [self.index_list])
-        window = full[idx]
-        if retain_index_list:
-            idx = tuple([slice(None)] * dim + [retain_index_list])
-            window[idx] = full[idx]
-        return window.to(device)
+        idx = [slice(None)] * dim + [self.index_list]
+        return full[idx].to(device)
 
     def add_window(self, full: torch.Tensor, to_add: torch.Tensor, dim=None) -> torch.Tensor:
         if dim is None:
             dim = self.dim
-        idx = tuple([slice(None)] * dim + [self.index_list])
+        idx = [slice(None)] * dim + [self.index_list]
         full[idx] += to_add
         return full
 
-    def get_region_index(self, num_regions: int) -> int:
-        region_idx = int(self.center_ratio * num_regions)
-        return min(max(region_idx, 0), num_regions - 1)
-
 
 class IndexListCallbacks:
     EVALUATE_CONTEXT_WINDOWS = "evaluate_context_windows"
     COMBINE_CONTEXT_WINDOW_RESULTS = "combine_context_window_results"
     EXECUTE_START = "execute_start"
     EXECUTE_CLEANUP = "execute_cleanup"
-    RESIZE_COND_ITEM = "resize_cond_item"
 
     def init_callbacks(self):
         return {}
@@ -105,8 +94,7 @@ class ContextFuseMethod:
 
 ContextResults = collections.namedtuple("ContextResults", ['window_idx', 'sub_conds_out', 'sub_conds', 'window'])
 class IndexListContextHandler(ContextHandlerABC):
-    def __init__(self, context_schedule: ContextSchedule, fuse_method: ContextFuseMethod, context_length: int=1, context_overlap: int=0, context_stride: int=1,
-                 closed_loop: bool=False, dim:int=0, freenoise: bool=False, cond_retain_index_list: list[int]=[], split_conds_to_windows: bool=False):
+    def __init__(self, context_schedule: ContextSchedule, fuse_method: ContextFuseMethod, context_length: int=1, context_overlap: int=0, context_stride: int=1, closed_loop=False, dim=0):
         self.context_schedule = context_schedule
         self.fuse_method = fuse_method
         self.context_length = context_length
@@ -115,18 +103,13 @@ class IndexListContextHandler(ContextHandlerABC):
         self.closed_loop = closed_loop
         self.dim = dim
         self._step = 0
-        self.freenoise = freenoise
-        self.cond_retain_index_list = [int(x.strip()) for x in cond_retain_index_list.split(",")] if cond_retain_index_list else []
-        self.split_conds_to_windows = split_conds_to_windows
 
         self.callbacks = {}
 
     def should_use_context(self, model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep: torch.Tensor, model_options: dict[str]) -> bool:
         # for now, assume first dim is batch - should have stored on BaseModel in actual implementation
         if x_in.size(self.dim) > self.context_length:
-            logging.info(f"Using context windows {self.context_length} with overlap {self.context_overlap} for {x_in.size(self.dim)} frames.")
-            if self.cond_retain_index_list:
-                logging.info(f"Retaining original cond for indexes: {self.cond_retain_index_list}")
+            logging.info(f"Using context windows {self.context_length} for {x_in.size(self.dim)} frames.")
             return True
         return False
 
@@ -140,11 +123,6 @@ class IndexListContextHandler(ContextHandlerABC):
             return None
         # reuse or resize cond items to match context requirements
         resized_cond = []
-        # if multiple conds, split based on primary region
-        if self.split_conds_to_windows and len(cond_in) > 1:
-            region = window.get_region_index(len(cond_in))
-            logging.info(f"Splitting conds to windows; using region {region} for window {window.index_list[0]}-{window.index_list[-1]} with center ratio {window.center_ratio:.3f}")
-            cond_in = [cond_in[region]]
         # cond object is a list containing a dict - outer list is irrelevant, so just loop through it
         for actual_cond in cond_in:
             resized_actual_cond = actual_cond.copy()
@@ -167,32 +145,13 @@ class IndexListContextHandler(ContextHandlerABC):
                         new_cond_item = cond_item.copy()
                         # when in dictionary, look for tensors and CONDCrossAttn [comfy/conds.py] (has cond attr that is a tensor)
                         for cond_key, cond_value in new_cond_item.items():
-                            # Allow callbacks to handle custom conditioning items
-                            handled = False
-                            for callback in comfy.patcher_extension.get_all_callbacks(
-                                IndexListCallbacks.RESIZE_COND_ITEM, self.callbacks
-                            ):
-                                result = callback(cond_key, cond_value, window, x_in, device, new_cond_item)
-                                if result is not None:
-                                    new_cond_item[cond_key] = result
-                                    handled = True
-                                    break
-                            if handled:
-                                continue
                             if isinstance(cond_value, torch.Tensor):
-                                if (self.dim < cond_value.ndim and cond_value(self.dim) == x_in.size(self.dim)) or \
-                                   (cond_value.ndim < self.dim and cond_value.size(0) == x_in.size(self.dim)):
+                                if cond_value.ndim < self.dim and cond_value.size(0) == x_in.size(self.dim):
                                     new_cond_item[cond_key] = window.get_tensor(cond_value, device)
-                            # Handle audio_embed (temporal dim is 1)
-                            elif cond_key == "audio_embed" and hasattr(cond_value, "cond") and isinstance(cond_value.cond, torch.Tensor):
-                                audio_cond = cond_value.cond
-                                if audio_cond.ndim > 1 and audio_cond.size(1) == x_in.size(self.dim):
-                                    new_cond_item[cond_key] = cond_value._copy_with(window.get_tensor(audio_cond, device, dim=1))
                             # if has cond that is a Tensor, check if needs to be subset
                             elif hasattr(cond_value, "cond") and isinstance(cond_value.cond, torch.Tensor):
-                                if  (self.dim < cond_value.cond.ndim and cond_value.cond.size(self.dim) == x_in.size(self.dim)) or \
-                                    (cond_value.cond.ndim < self.dim and cond_value.cond.size(0) == x_in.size(self.dim)):
-                                    new_cond_item[cond_key] = cond_value._copy_with(window.get_tensor(cond_value.cond, device, retain_index_list=self.cond_retain_index_list))
+                                if cond_value.cond.ndim < self.dim and cond_value.cond.size(0) == x_in.size(self.dim):
+                                    new_cond_item[cond_key] = cond_value._copy_with(window.get_tensor(cond_value.cond, device))
                             elif cond_key == "num_video_frames": # for SVD
                                 new_cond_item[cond_key] = cond_value._copy_with(cond_value.cond)
                                 new_cond_item[cond_key].cond = window.context_length
@@ -205,7 +164,7 @@ class IndexListContextHandler(ContextHandlerABC):
         return resized_cond
 
     def set_step(self, timestep: torch.Tensor, model_options: dict[str]):
-        mask = torch.isclose(model_options["transformer_options"]["sample_sigmas"], timestep[0], rtol=0.0001)
+        mask = torch.isclose(model_options["transformer_options"]["sample_sigmas"], timestep, rtol=0.0001)
         matches = torch.nonzero(mask)
         if torch.numel(matches) == 0:
             raise Exception("No sample_sigmas matched current timestep; something went wrong.")
@@ -214,7 +173,7 @@ class IndexListContextHandler(ContextHandlerABC):
     def get_context_windows(self, model: BaseModel, x_in: torch.Tensor, model_options: dict[str]) -> list[IndexListContextWindow]:
         full_length = x_in.size(self.dim) # TODO: choose dim based on model
         context_windows = self.context_schedule.func(full_length, self, model_options)
-        context_windows = [IndexListContextWindow(window, dim=self.dim, total_frames=full_length) for window in context_windows]
+        context_windows = [IndexListContextWindow(window, dim=self.dim) for window in context_windows]
         return context_windows
 
     def execute(self, calc_cond_batch: Callable, model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep: torch.Tensor, model_options: dict[str]):
@@ -291,8 +250,8 @@ class IndexListContextHandler(ContextHandlerABC):
                     prev_weight = (bias_total / (bias_total + bias))
                     new_weight = (bias / (bias_total + bias))
                     # account for dims of tensors
-                    idx_window = tuple([slice(None)] * self.dim + [idx])
-                    pos_window = tuple([slice(None)] * self.dim + [pos])
+                    idx_window = [slice(None)] * self.dim + [idx]
+                    pos_window = [slice(None)] * self.dim + [pos]
                     # apply new values
                     conds_final[i][idx_window] = conds_final[i][idx_window] * prev_weight + sub_conds_out[i][pos_window] * new_weight
                     biases_final[i][idx] = bias_total + bias
@@ -328,28 +287,6 @@ def create_prepare_sampling_wrapper(model: ModelPatcher):
     )
 
 
-def _sampler_sample_wrapper(executor, guider, sigmas, extra_args, callback, noise, *args, **kwargs):
-    model_options = extra_args.get("model_options", None)
-    if model_options is None:
-        raise Exception("model_options not found in sampler_sample_wrapper; this should never happen, something went wrong.")
-    handler: IndexListContextHandler = model_options.get("context_handler", None)
-    if handler is None:
-        raise Exception("context_handler not found in sampler_sample_wrapper; this should never happen, something went wrong.")
-    if not handler.freenoise:
-        return executor(guider, sigmas, extra_args, callback, noise, *args, **kwargs)
-    noise = apply_freenoise(noise, handler.dim, handler.context_length, handler.context_overlap, extra_args["seed"])
-
-    return executor(guider, sigmas, extra_args, callback, noise, *args, **kwargs)
-
-
-def create_sampler_sample_wrapper(model: ModelPatcher):
-    model.add_wrapper_with_key(
-        comfy.patcher_extension.WrappersMP.SAMPLER_SAMPLE,
-        "ContextWindows_sampler_sample",
-        _sampler_sample_wrapper
-    )
-
-
 def match_weights_to_dim(weights: list[float], x_in: torch.Tensor, dim: int, device=None) -> torch.Tensor:
     total_dims = len(x_in.shape)
     weights_tensor = torch.Tensor(weights).to(device=device)
@@ -601,29 +538,3 @@ def shift_window_to_end(window: list[int], num_frames: int):
     for i in range(len(window)):
         # 2) add end_delta to each val to slide windows to end
         window[i] = window[i] + end_delta
-
-
-# https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved/blob/90fb1331201a4b29488089e4fbffc0d82cc6d0a9/animatediff/sample_settings.py#L465
-def apply_freenoise(noise: torch.Tensor, dim: int, context_length: int, context_overlap: int, seed: int):
-    logging.info("Context windows: Applying FreeNoise")
-    generator = torch.Generator(device='cpu').manual_seed(seed)
-    latent_video_length = noise.shape[dim]
-    delta = context_length - context_overlap
-
-    for start_idx in range(0, latent_video_length - context_length, delta):
-        place_idx = start_idx + context_length
-
-        actual_delta = min(delta, latent_video_length - place_idx)
-        if actual_delta <= 0:
-            break
-
-        list_idx = torch.randperm(actual_delta, generator=generator, device='cpu') + start_idx
-
-        source_slice = [slice(None)] * noise.ndim
-        source_slice[dim] = list_idx
-        target_slice = [slice(None)] * noise.ndim
-        target_slice[dim] = slice(place_idx, place_idx + actual_delta)
-
-        noise[tuple(target_slice)] = noise[tuple(source_slice)]
-
-    return noise

comfy/k_diffusion/sampling.py

@@ -1557,13 +1557,10 @@ def sample_er_sde(model, x, sigmas, extra_args=None, callback=None, disable=None
 
 
 @torch.no_grad()
-def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=0.5, solver_type="phi_1"):
+def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=0.5):
     """SEEDS-2 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 2.
     arXiv: https://arxiv.org/abs/2305.14267 (NeurIPS 2023)
     """
-    if solver_type not in {"phi_1", "phi_2"}:
-        raise ValueError("solver_type must be 'phi_1' or 'phi_2'")
-
     extra_args = {} if extra_args is None else extra_args
     seed = extra_args.get("seed", None)
     noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
@@ -1603,14 +1600,8 @@ def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=Non
         denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
 
         # Step 2
-        if solver_type == "phi_1":
-            denoised_d = torch.lerp(denoised, denoised_2, fac)
-            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * ei_h_phi_1(-h_eta) * denoised_d
-        elif solver_type == "phi_2":
-            b2 = ei_h_phi_2(-h_eta) / r
-            b1 = ei_h_phi_1(-h_eta) - b2
-            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * (b1 * denoised + b2 * denoised_2)
-
+        denoised_d = torch.lerp(denoised, denoised_2, fac)
+        x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * ei_h_phi_1(-h_eta) * denoised_d
         if inject_noise:
             segment_factor = (r - 1) * h * eta
             sde_noise = sde_noise * segment_factor.exp()
@@ -1618,17 +1609,6 @@ def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=Non
             x = x + sde_noise * sigmas[i + 1] * s_noise
     return x
 
-@torch.no_grad()
-def sample_exp_heun_2_x0(model, x, sigmas, extra_args=None, callback=None, disable=None, solver_type="phi_2"):
-    """Deterministic exponential Heun second order method in data prediction (x0) and logSNR time."""
-    return sample_seeds_2(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=0.0, s_noise=0.0, noise_sampler=None, r=1.0, solver_type=solver_type)
-
-
-@torch.no_grad()
-def sample_exp_heun_2_x0_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type="phi_2"):
-    """Stochastic exponential Heun second order method in data prediction (x0) and logSNR time."""
-    return sample_seeds_2(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, r=1.0, solver_type=solver_type)
-
 
 @torch.no_grad()
 def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r_1=1./3, r_2=2./3):
@@ -1776,7 +1756,7 @@ def sample_sa_solver(model, x, sigmas, extra_args=None, callback=None, disable=F
         # Predictor
         if sigmas[i + 1] == 0:
             # Denoising step
-            x_pred = denoised
+            x = denoised
         else:
             tau_t = tau_func(sigmas[i + 1])
             curr_lambdas = lambdas[i - predictor_order_used + 1:i + 1]
@@ -1797,7 +1777,7 @@ def sample_sa_solver(model, x, sigmas, extra_args=None, callback=None, disable=F
             if tau_t > 0 and s_noise > 0:
                 noise = noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * tau_t ** 2 * h).expm1().neg().sqrt() * s_noise
                 x_pred = x_pred + noise
-    return x_pred
+    return x
 
 
 @torch.no_grad()

comfy/latent_formats.py

@@ -6,7 +6,6 @@ class LatentFormat:
     latent_dimensions = 2
     latent_rgb_factors = None
     latent_rgb_factors_bias = None
-    latent_rgb_factors_reshape = None
     taesd_decoder_name = None
 
     def process_in(self, latent):
@@ -179,54 +178,6 @@ class Flux(SD3):
     def process_out(self, latent):
         return (latent / self.scale_factor) + self.shift_factor
 
-class Flux2(LatentFormat):
-    latent_channels = 128
-
-    def __init__(self):
-        self.latent_rgb_factors =[
-            [0.0058, 0.0113, 0.0073],
-            [0.0495, 0.0443, 0.0836],
-            [-0.0099, 0.0096, 0.0644],
-            [0.2144, 0.3009, 0.3652],
-            [0.0166, -0.0039, -0.0054],
-            [0.0157, 0.0103, -0.0160],
-            [-0.0398, 0.0902, -0.0235],
-            [-0.0052, 0.0095, 0.0109],
-            [-0.3527, -0.2712, -0.1666],
-            [-0.0301, -0.0356, -0.0180],
-            [-0.0107, 0.0078, 0.0013],
-            [0.0746, 0.0090, -0.0941],
-            [0.0156, 0.0169, 0.0070],
-            [-0.0034, -0.0040, -0.0114],
-            [0.0032, 0.0181, 0.0080],
-            [-0.0939, -0.0008, 0.0186],
-            [0.0018, 0.0043, 0.0104],
-            [0.0284, 0.0056, -0.0127],
-            [-0.0024, -0.0022, -0.0030],
-            [0.1207, -0.0026, 0.0065],
-            [0.0128, 0.0101, 0.0142],
-            [0.0137, -0.0072, -0.0007],
-            [0.0095, 0.0092, -0.0059],
-            [0.0000, -0.0077, -0.0049],
-            [-0.0465, -0.0204, -0.0312],
-            [0.0095, 0.0012, -0.0066],
-            [0.0290, -0.0034, 0.0025],
-            [0.0220, 0.0169, -0.0048],
-            [-0.0332, -0.0457, -0.0468],
-            [-0.0085, 0.0389, 0.0609],
-            [-0.0076, 0.0003, -0.0043],
-            [-0.0111, -0.0460, -0.0614],
-        ]
-
-        self.latent_rgb_factors_bias = [-0.0329, -0.0718, -0.0851]
-        self.latent_rgb_factors_reshape = lambda t: t.reshape(t.shape[0], 32, 2, 2, t.shape[-2], t.shape[-1]).permute(0, 1, 4, 2, 5, 3).reshape(t.shape[0], 32, t.shape[-2] * 2, t.shape[-1] * 2)
-
-    def process_in(self, latent):
-        return latent
-
-    def process_out(self, latent):
-        return latent
-
 class Mochi(LatentFormat):
     latent_channels = 12
     latent_dimensions = 3
@@ -431,7 +382,6 @@ class HunyuanVideo(LatentFormat):
     ]
 
     latent_rgb_factors_bias = [ 0.0259, -0.0192, -0.0761]
-    taesd_decoder_name = "taehv"
 
 class Cosmos1CV8x8x8(LatentFormat):
     latent_channels = 16
@@ -495,7 +445,7 @@ class Wan21(LatentFormat):
         ]).view(1, self.latent_channels, 1, 1, 1)
 
 
-        self.taesd_decoder_name = "lighttaew2_1"
+        self.taesd_decoder_name = None #TODO
 
     def process_in(self, latent):
         latents_mean = self.latents_mean.to(latent.device, latent.dtype)
@@ -566,7 +516,6 @@ class Wan22(Wan21):
 
     def __init__(self):
         self.scale_factor = 1.0
-        self.taesd_decoder_name = "lighttaew2_2"
         self.latents_mean = torch.tensor([
                 -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557,
                 -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825,
@@ -662,67 +611,6 @@ class HunyuanImage21Refiner(LatentFormat):
     latent_dimensions = 3
     scale_factor = 1.03682
 
-    def process_in(self, latent):
-        out = latent * self.scale_factor
-        out = torch.cat((out[:, :, :1], out), dim=2)
-        out = out.permute(0, 2, 1, 3, 4)
-        b, f_times_2, c, h, w = out.shape
-        out = out.reshape(b, f_times_2 // 2, 2 * c, h, w)
-        out = out.permute(0, 2, 1, 3, 4).contiguous()
-        return out
-
-    def process_out(self, latent):
-        z = latent / self.scale_factor
-        z = z.permute(0, 2, 1, 3, 4)
-        b, f, c, h, w = z.shape
-        z = z.reshape(b, f, 2, c // 2, h, w)
-        z = z.permute(0, 1, 2, 3, 4, 5).reshape(b, f * 2, c // 2, h, w)
-        z = z.permute(0, 2, 1, 3, 4)
-        z = z[:, :, 1:]
-        return z
-
-class HunyuanVideo15(LatentFormat):
-    latent_rgb_factors = [
-        [ 0.0568, -0.0521, -0.0131],
-        [ 0.0014,  0.0735,  0.0326],
-        [ 0.0186,  0.0531, -0.0138],
-        [-0.0031,  0.0051,  0.0288],
-        [ 0.0110,  0.0556,  0.0432],
-        [-0.0041, -0.0023, -0.0485],
-        [ 0.0530,  0.0413,  0.0253],
-        [ 0.0283,  0.0251,  0.0339],
-        [ 0.0277, -0.0372, -0.0093],
-        [ 0.0393,  0.0944,  0.1131],
-        [ 0.0020,  0.0251,  0.0037],
-        [-0.0017,  0.0012,  0.0234],
-        [ 0.0468,  0.0436,  0.0203],
-        [ 0.0354,  0.0439, -0.0233],
-        [ 0.0090,  0.0123,  0.0346],
-        [ 0.0382,  0.0029,  0.0217],
-        [ 0.0261, -0.0300,  0.0030],
-        [-0.0088, -0.0220, -0.0283],
-        [-0.0272, -0.0121, -0.0363],
-        [-0.0664, -0.0622,  0.0144],
-        [ 0.0414,  0.0479,  0.0529],
-        [ 0.0355,  0.0612, -0.0247],
-        [ 0.0147,  0.0264,  0.0174],
-        [ 0.0438,  0.0038,  0.0542],
-        [ 0.0431, -0.0573, -0.0033],
-        [-0.0162, -0.0211, -0.0406],
-        [-0.0487, -0.0295, -0.0393],
-        [ 0.0005, -0.0109,  0.0253],
-        [ 0.0296,  0.0591,  0.0353],
-        [ 0.0119,  0.0181, -0.0306],
-        [-0.0085, -0.0362,  0.0229],
-        [ 0.0005, -0.0106,  0.0242]
-    ]
-
-    latent_rgb_factors_bias = [ 0.0456, -0.0202, -0.0644]
-    latent_channels = 32
-    latent_dimensions = 3
-    scale_factor = 1.03682
-    taesd_decoder_name = "lighttaehy1_5"
-
 class Hunyuan3Dv2(LatentFormat):
     latent_channels = 64
     latent_dimensions = 1

comfy/ldm/chroma/layers.py

@@ -1,15 +1,15 @@
 import torch
 from torch import Tensor, nn
 
+from comfy.ldm.flux.math import attention
 from comfy.ldm.flux.layers import (
     MLPEmbedder,
     RMSNorm,
+    QKNorm,
+    SelfAttention,
     ModulationOut,
 )
 
-# TODO: remove this in a few months
-SingleStreamBlock = None
-DoubleStreamBlock = None
 
 
 class ChromaModulationOut(ModulationOut):
@@ -48,6 +48,124 @@ class Approximator(nn.Module):
         return x
 
 
+class DoubleStreamBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, flipped_img_txt=False, dtype=None, device=None, operations=None):
+        super().__init__()
+
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.img_norm1 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, dtype=dtype, device=device, operations=operations)
+
+        self.img_norm2 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.img_mlp = nn.Sequential(
+            operations.Linear(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, device=device),
+            nn.GELU(approximate="tanh"),
+            operations.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
+        )
+
+        self.txt_norm1 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, dtype=dtype, device=device, operations=operations)
+
+        self.txt_norm2 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.txt_mlp = nn.Sequential(
+            operations.Linear(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, device=device),
+            nn.GELU(approximate="tanh"),
+            operations.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
+        )
+        self.flipped_img_txt = flipped_img_txt
+
+    def forward(self, img: Tensor, txt: Tensor, pe: Tensor, vec: Tensor, attn_mask=None, transformer_options={}):
+        (img_mod1, img_mod2), (txt_mod1, txt_mod2) = vec
+
+        # prepare image for attention
+        img_modulated = torch.addcmul(img_mod1.shift, 1 + img_mod1.scale, self.img_norm1(img))
+        img_qkv = self.img_attn.qkv(img_modulated)
+        img_q, img_k, img_v = img_qkv.view(img_qkv.shape[0], img_qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+
+        # prepare txt for attention
+        txt_modulated = torch.addcmul(txt_mod1.shift, 1 + txt_mod1.scale, self.txt_norm1(txt))
+        txt_qkv = self.txt_attn.qkv(txt_modulated)
+        txt_q, txt_k, txt_v = txt_qkv.view(txt_qkv.shape[0], txt_qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+        # run actual attention
+        attn = attention(torch.cat((txt_q, img_q), dim=2),
+                         torch.cat((txt_k, img_k), dim=2),
+                         torch.cat((txt_v, img_v), dim=2),
+                         pe=pe, mask=attn_mask, transformer_options=transformer_options)
+
+        txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+        # calculate the img bloks
+        img.addcmul_(img_mod1.gate, self.img_attn.proj(img_attn))
+        img.addcmul_(img_mod2.gate, self.img_mlp(torch.addcmul(img_mod2.shift, 1 + img_mod2.scale, self.img_norm2(img))))
+
+        # calculate the txt bloks
+        txt.addcmul_(txt_mod1.gate, self.txt_attn.proj(txt_attn))
+        txt.addcmul_(txt_mod2.gate, self.txt_mlp(torch.addcmul(txt_mod2.shift, 1 + txt_mod2.scale, self.txt_norm2(txt))))
+
+        if txt.dtype == torch.float16:
+            txt = torch.nan_to_num(txt, nan=0.0, posinf=65504, neginf=-65504)
+
+        return img, txt
+
+
+class SingleStreamBlock(nn.Module):
+    """
+    A DiT block with parallel linear layers as described in
+    https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qk_scale: float = None,
+        dtype=None,
+        device=None,
+        operations=None
+    ):
+        super().__init__()
+        self.hidden_dim = hidden_size
+        self.num_heads = num_heads
+        head_dim = hidden_size // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        # qkv and mlp_in
+        self.linear1 = operations.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim, dtype=dtype, device=device)
+        # proj and mlp_out
+        self.linear2 = operations.Linear(hidden_size + self.mlp_hidden_dim, hidden_size, dtype=dtype, device=device)
+
+        self.norm = QKNorm(head_dim, dtype=dtype, device=device, operations=operations)
+
+        self.hidden_size = hidden_size
+        self.pre_norm = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+
+        self.mlp_act = nn.GELU(approximate="tanh")
+
+    def forward(self, x: Tensor, pe: Tensor, vec: Tensor, attn_mask=None, transformer_options={}) -> Tensor:
+        mod = vec
+        x_mod = torch.addcmul(mod.shift, 1 + mod.scale, self.pre_norm(x))
+        qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+
+        q, k, v = qkv.view(qkv.shape[0], qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        q, k = self.norm(q, k, v)
+
+        # compute attention
+        attn = attention(q, k, v, pe=pe, mask=attn_mask, transformer_options=transformer_options)
+        # compute activation in mlp stream, cat again and run second linear layer
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+        x.addcmul_(mod.gate, output)
+        if x.dtype == torch.float16:
+            x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
+        return x
+
+
 class LastLayer(nn.Module):
     def __init__(self, hidden_size: int, patch_size: int, out_channels: int, dtype=None, device=None, operations=None):
         super().__init__()

comfy/ldm/chroma/model.py

@@ -11,12 +11,12 @@ import comfy.ldm.common_dit
 from comfy.ldm.flux.layers import (
     EmbedND,
     timestep_embedding,
-    DoubleStreamBlock,
-    SingleStreamBlock,
 )
 
 from .layers import (
+    DoubleStreamBlock,
     LastLayer,
+    SingleStreamBlock,
     Approximator,
     ChromaModulationOut,
 )
@@ -40,8 +40,7 @@ class ChromaParams:
     out_dim: int
     hidden_dim: int
     n_layers: int
-    txt_ids_dims: list
-    vec_in_dim: int
+
 
 
 
@@ -91,7 +90,6 @@ class Chroma(nn.Module):
                     self.num_heads,
                     mlp_ratio=params.mlp_ratio,
                     qkv_bias=params.qkv_bias,
-                    modulation=False,
                     dtype=dtype, device=device, operations=operations
                 )
                 for _ in range(params.depth)
@@ -100,7 +98,7 @@ class Chroma(nn.Module):
 
         self.single_blocks = nn.ModuleList(
             [
-                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio, modulation=False, dtype=dtype, device=device, operations=operations)
+                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio, dtype=dtype, device=device, operations=operations)
                 for _ in range(params.depth_single_blocks)
             ]
         )
@@ -180,10 +178,7 @@ class Chroma(nn.Module):
         pe = self.pe_embedder(ids)
 
         blocks_replace = patches_replace.get("dit", {})
-        transformer_options["total_blocks"] = len(self.double_blocks)
-        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.double_blocks):
-            transformer_options["block_index"] = i
             if i not in self.skip_mmdit:
                 double_mod = (
                     self.get_modulations(mod_vectors, "double_img", idx=i),
@@ -226,10 +221,7 @@ class Chroma(nn.Module):
 
         img = torch.cat((txt, img), 1)
 
-        transformer_options["total_blocks"] = len(self.single_blocks)
-        transformer_options["block_type"] = "single"
         for i, block in enumerate(self.single_blocks):
-            transformer_options["block_index"] = i
             if i not in self.skip_dit:
                 single_mod = self.get_modulations(mod_vectors, "single", idx=i)
                 if ("single_block", i) in blocks_replace:

comfy/ldm/chroma_radiance/model.py

@@ -10,10 +10,12 @@ from torch import Tensor, nn
 from einops import repeat
 import comfy.ldm.common_dit
 
-from comfy.ldm.flux.layers import EmbedND, DoubleStreamBlock, SingleStreamBlock
+from comfy.ldm.flux.layers import EmbedND
 
 from comfy.ldm.chroma.model import Chroma, ChromaParams
 from comfy.ldm.chroma.layers import (
+    DoubleStreamBlock,
+    SingleStreamBlock,
     Approximator,
 )
 from .layers import (
@@ -37,7 +39,7 @@ class ChromaRadianceParams(ChromaParams):
     nerf_final_head_type: str
     # None means use the same dtype as the model.
     nerf_embedder_dtype: Optional[torch.dtype]
-    use_x0: bool
+
 
 class ChromaRadiance(Chroma):
     """
@@ -87,6 +89,7 @@ class ChromaRadiance(Chroma):
                     dtype=dtype, device=device, operations=operations
                 )
 
+
         self.double_blocks = nn.ModuleList(
             [
                 DoubleStreamBlock(
@@ -94,7 +97,6 @@ class ChromaRadiance(Chroma):
                     self.num_heads,
                     mlp_ratio=params.mlp_ratio,
                     qkv_bias=params.qkv_bias,
-                    modulation=False,
                     dtype=dtype, device=device, operations=operations
                 )
                 for _ in range(params.depth)
@@ -107,7 +109,6 @@ class ChromaRadiance(Chroma):
                     self.hidden_size,
                     self.num_heads,
                     mlp_ratio=params.mlp_ratio,
-                    modulation=False,
                     dtype=dtype, device=device, operations=operations,
                 )
                 for _ in range(params.depth_single_blocks)
@@ -159,9 +160,6 @@ class ChromaRadiance(Chroma):
         self.skip_dit = []
         self.lite = False
 
-        if params.use_x0:
-            self.register_buffer("__x0__", torch.tensor([]))
-
     @property
     def _nerf_final_layer(self) -> nn.Module:
         if self.params.nerf_final_head_type == "linear":
@@ -279,12 +277,6 @@ class ChromaRadiance(Chroma):
         params_dict |= overrides
         return params.__class__(**params_dict)
 
-    def _apply_x0_residual(self, predicted, noisy, timesteps):
-
-        # non zero during training to prevent 0 div
-        eps = 0.0
-        return (noisy - predicted) / (timesteps.view(-1,1,1,1) + eps)
-
     def _forward(
         self,
         x: Tensor,
@@ -325,11 +317,4 @@ class ChromaRadiance(Chroma):
             transformer_options,
             attn_mask=kwargs.get("attention_mask", None),
         )
-
-        out = self.forward_nerf(img, img_out, params)[:, :, :h, :w]
-
-        # If x0 variant → v-pred, just return this instead
-        if hasattr(self, "__x0__"):
-            out = self._apply_x0_residual(out, img, timestep)
-        return out
-
+        return self.forward_nerf(img, img_out, params)[:, :, :h, :w]

comfy/ldm/flux/layers.py

@@ -48,44 +48,15 @@ def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 10
     return embedding
 
 class MLPEmbedder(nn.Module):
-    def __init__(self, in_dim: int, hidden_dim: int, bias=True, dtype=None, device=None, operations=None):
+    def __init__(self, in_dim: int, hidden_dim: int, dtype=None, device=None, operations=None):
         super().__init__()
-        self.in_layer = operations.Linear(in_dim, hidden_dim, bias=bias, dtype=dtype, device=device)
+        self.in_layer = operations.Linear(in_dim, hidden_dim, bias=True, dtype=dtype, device=device)
         self.silu = nn.SiLU()
-        self.out_layer = operations.Linear(hidden_dim, hidden_dim, bias=bias, dtype=dtype, device=device)
+        self.out_layer = operations.Linear(hidden_dim, hidden_dim, bias=True, dtype=dtype, device=device)
 
     def forward(self, x: Tensor) -> Tensor:
         return self.out_layer(self.silu(self.in_layer(x)))
 
-class YakMLP(nn.Module):
-    def __init__(self, hidden_size: int, intermediate_size: int, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.gate_proj = operations.Linear(self.hidden_size, self.intermediate_size, bias=True, dtype=dtype, device=device)
-        self.up_proj = operations.Linear(self.hidden_size, self.intermediate_size, bias=True, dtype=dtype, device=device)
-        self.down_proj = operations.Linear(self.intermediate_size, self.hidden_size, bias=True, dtype=dtype, device=device)
-        self.act_fn = nn.SiLU()
-
-    def forward(self, x: Tensor) -> Tensor:
-        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-        return down_proj
-
-def build_mlp(hidden_size, mlp_hidden_dim, mlp_silu_act=False, yak_mlp=False, dtype=None, device=None, operations=None):
-    if yak_mlp:
-        return YakMLP(hidden_size, mlp_hidden_dim, dtype=dtype, device=device, operations=operations)
-    if mlp_silu_act:
-        return nn.Sequential(
-            operations.Linear(hidden_size, mlp_hidden_dim * 2, bias=False, dtype=dtype, device=device),
-            SiLUActivation(),
-            operations.Linear(mlp_hidden_dim, hidden_size, bias=False, dtype=dtype, device=device),
-        )
-    else:
-        return nn.Sequential(
-            operations.Linear(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, device=device),
-            nn.GELU(approximate="tanh"),
-            operations.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
-        )
 
 class RMSNorm(torch.nn.Module):
     def __init__(self, dim: int, dtype=None, device=None, operations=None):
@@ -109,14 +80,14 @@ class QKNorm(torch.nn.Module):
 
 
 class SelfAttention(nn.Module):
-    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False, proj_bias: bool = True, dtype=None, device=None, operations=None):
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False, dtype=None, device=None, operations=None):
         super().__init__()
         self.num_heads = num_heads
         head_dim = dim // num_heads
 
         self.qkv = operations.Linear(dim, dim * 3, bias=qkv_bias, dtype=dtype, device=device)
         self.norm = QKNorm(head_dim, dtype=dtype, device=device, operations=operations)
-        self.proj = operations.Linear(dim, dim, bias=proj_bias, dtype=dtype, device=device)
+        self.proj = operations.Linear(dim, dim, dtype=dtype, device=device)
 
 
 @dataclass
@@ -127,11 +98,11 @@ class ModulationOut:
 
 
 class Modulation(nn.Module):
-    def __init__(self, dim: int, double: bool, bias=True, dtype=None, device=None, operations=None):
+    def __init__(self, dim: int, double: bool, dtype=None, device=None, operations=None):
         super().__init__()
         self.is_double = double
         self.multiplier = 6 if double else 3
-        self.lin = operations.Linear(dim, self.multiplier * dim, bias=bias, dtype=dtype, device=device)
+        self.lin = operations.Linear(dim, self.multiplier * dim, bias=True, dtype=dtype, device=device)
 
     def forward(self, vec: Tensor) -> tuple:
         if vec.ndim == 2:
@@ -158,107 +129,77 @@ def apply_mod(tensor, m_mult, m_add=None, modulation_dims=None):
         return tensor
 
 
-class SiLUActivation(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.gate_fn = nn.SiLU()
-
-    def forward(self, x: Tensor) -> Tensor:
-        x1, x2 = x.chunk(2, dim=-1)
-        return self.gate_fn(x1) * x2
-
-
 class DoubleStreamBlock(nn.Module):
-    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, flipped_img_txt=False, modulation=True, mlp_silu_act=False, proj_bias=True, yak_mlp=False, dtype=None, device=None, operations=None):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, flipped_img_txt=False, dtype=None, device=None, operations=None):
         super().__init__()
 
         mlp_hidden_dim = int(hidden_size * mlp_ratio)
         self.num_heads = num_heads
         self.hidden_size = hidden_size
-        self.modulation = modulation
-
-        if self.modulation:
-            self.img_mod = Modulation(hidden_size, double=True, dtype=dtype, device=device, operations=operations)
-
+        self.img_mod = Modulation(hidden_size, double=True, dtype=dtype, device=device, operations=operations)
         self.img_norm1 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, proj_bias=proj_bias, dtype=dtype, device=device, operations=operations)
+        self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, dtype=dtype, device=device, operations=operations)
 
         self.img_norm2 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.img_mlp = nn.Sequential(
+            operations.Linear(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, device=device),
+            nn.GELU(approximate="tanh"),
+            operations.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
+        )
 
-        self.img_mlp = build_mlp(hidden_size, mlp_hidden_dim, mlp_silu_act=mlp_silu_act, yak_mlp=yak_mlp, dtype=dtype, device=device, operations=operations)
-
-        if self.modulation:
-            self.txt_mod = Modulation(hidden_size, double=True, dtype=dtype, device=device, operations=operations)
-
+        self.txt_mod = Modulation(hidden_size, double=True, dtype=dtype, device=device, operations=operations)
         self.txt_norm1 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, proj_bias=proj_bias, dtype=dtype, device=device, operations=operations)
+        self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, dtype=dtype, device=device, operations=operations)
 
         self.txt_norm2 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-
-        self.txt_mlp = build_mlp(hidden_size, mlp_hidden_dim, mlp_silu_act=mlp_silu_act, yak_mlp=yak_mlp, dtype=dtype, device=device, operations=operations)
-
+        self.txt_mlp = nn.Sequential(
+            operations.Linear(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, device=device),
+            nn.GELU(approximate="tanh"),
+            operations.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
+        )
         self.flipped_img_txt = flipped_img_txt
 
     def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, modulation_dims_img=None, modulation_dims_txt=None, transformer_options={}):
-        if self.modulation:
-            img_mod1, img_mod2 = self.img_mod(vec)
-            txt_mod1, txt_mod2 = self.txt_mod(vec)
-        else:
-            (img_mod1, img_mod2), (txt_mod1, txt_mod2) = vec
+        img_mod1, img_mod2 = self.img_mod(vec)
+        txt_mod1, txt_mod2 = self.txt_mod(vec)
 
         # prepare image for attention
         img_modulated = self.img_norm1(img)
         img_modulated = apply_mod(img_modulated, (1 + img_mod1.scale), img_mod1.shift, modulation_dims_img)
         img_qkv = self.img_attn.qkv(img_modulated)
-        del img_modulated
         img_q, img_k, img_v = img_qkv.view(img_qkv.shape[0], img_qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
-        del img_qkv
         img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
 
         # prepare txt for attention
         txt_modulated = self.txt_norm1(txt)
         txt_modulated = apply_mod(txt_modulated, (1 + txt_mod1.scale), txt_mod1.shift, modulation_dims_txt)
         txt_qkv = self.txt_attn.qkv(txt_modulated)
-        del txt_modulated
         txt_q, txt_k, txt_v = txt_qkv.view(txt_qkv.shape[0], txt_qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
-        del txt_qkv
         txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
 
         if self.flipped_img_txt:
-            q = torch.cat((img_q, txt_q), dim=2)
-            del img_q, txt_q
-            k = torch.cat((img_k, txt_k), dim=2)
-            del img_k, txt_k
-            v = torch.cat((img_v, txt_v), dim=2)
-            del img_v, txt_v
             # run actual attention
-            attn = attention(q, k, v,
+            attn = attention(torch.cat((img_q, txt_q), dim=2),
+                             torch.cat((img_k, txt_k), dim=2),
+                             torch.cat((img_v, txt_v), dim=2),
                              pe=pe, mask=attn_mask, transformer_options=transformer_options)
-            del q, k, v
 
             img_attn, txt_attn = attn[:, : img.shape[1]], attn[:, img.shape[1]:]
         else:
-            q = torch.cat((txt_q, img_q), dim=2)
-            del txt_q, img_q
-            k = torch.cat((txt_k, img_k), dim=2)
-            del txt_k, img_k
-            v = torch.cat((txt_v, img_v), dim=2)
-            del txt_v, img_v
             # run actual attention
-            attn = attention(q, k, v,
+            attn = attention(torch.cat((txt_q, img_q), dim=2),
+                             torch.cat((txt_k, img_k), dim=2),
+                             torch.cat((txt_v, img_v), dim=2),
                              pe=pe, mask=attn_mask, transformer_options=transformer_options)
-            del q, k, v
 
             txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1]:]
 
         # calculate the img bloks
         img += apply_mod(self.img_attn.proj(img_attn), img_mod1.gate, None, modulation_dims_img)
-        del img_attn
         img += apply_mod(self.img_mlp(apply_mod(self.img_norm2(img), (1 + img_mod2.scale), img_mod2.shift, modulation_dims_img)), img_mod2.gate, None, modulation_dims_img)
 
         # calculate the txt bloks
         txt += apply_mod(self.txt_attn.proj(txt_attn), txt_mod1.gate, None, modulation_dims_txt)
-        del txt_attn
         txt += apply_mod(self.txt_mlp(apply_mod(self.txt_norm2(txt), (1 + txt_mod2.scale), txt_mod2.shift, modulation_dims_txt)), txt_mod2.gate, None, modulation_dims_txt)
 
         if txt.dtype == torch.float16:
@@ -279,10 +220,6 @@ class SingleStreamBlock(nn.Module):
         num_heads: int,
         mlp_ratio: float = 4.0,
         qk_scale: float = None,
-        modulation=True,
-        mlp_silu_act=False,
-        bias=True,
-        yak_mlp=False,
         dtype=None,
         device=None,
         operations=None
@@ -294,55 +231,30 @@ class SingleStreamBlock(nn.Module):
         self.scale = qk_scale or head_dim**-0.5
 
         self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
-
-        self.mlp_hidden_dim_first = self.mlp_hidden_dim
-        self.yak_mlp = yak_mlp
-        if mlp_silu_act:
-            self.mlp_hidden_dim_first = int(hidden_size * mlp_ratio * 2)
-            self.mlp_act = SiLUActivation()
-        else:
-            self.mlp_act = nn.GELU(approximate="tanh")
-
-        if self.yak_mlp:
-            self.mlp_hidden_dim_first *= 2
-            self.mlp_act = nn.SiLU()
-
         # qkv and mlp_in
-        self.linear1 = operations.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim_first, bias=bias, dtype=dtype, device=device)
+        self.linear1 = operations.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim, dtype=dtype, device=device)
         # proj and mlp_out
-        self.linear2 = operations.Linear(hidden_size + self.mlp_hidden_dim, hidden_size, bias=bias, dtype=dtype, device=device)
+        self.linear2 = operations.Linear(hidden_size + self.mlp_hidden_dim, hidden_size, dtype=dtype, device=device)
 
         self.norm = QKNorm(head_dim, dtype=dtype, device=device, operations=operations)
 
         self.hidden_size = hidden_size
         self.pre_norm = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
 
-        if modulation:
-            self.modulation = Modulation(hidden_size, double=False, dtype=dtype, device=device, operations=operations)
-        else:
-            self.modulation = None
+        self.mlp_act = nn.GELU(approximate="tanh")
+        self.modulation = Modulation(hidden_size, double=False, dtype=dtype, device=device, operations=operations)
 
     def forward(self, x: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, modulation_dims=None, transformer_options={}) -> Tensor:
-        if self.modulation:
-            mod, _ = self.modulation(vec)
-        else:
-            mod = vec
-
-        qkv, mlp = torch.split(self.linear1(apply_mod(self.pre_norm(x), (1 + mod.scale), mod.shift, modulation_dims)), [3 * self.hidden_size, self.mlp_hidden_dim_first], dim=-1)
+        mod, _ = self.modulation(vec)
+        qkv, mlp = torch.split(self.linear1(apply_mod(self.pre_norm(x), (1 + mod.scale), mod.shift, modulation_dims)), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
 
         q, k, v = qkv.view(qkv.shape[0], qkv.shape[1], 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
-        del qkv
         q, k = self.norm(q, k, v)
 
         # compute attention
         attn = attention(q, k, v, pe=pe, mask=attn_mask, transformer_options=transformer_options)
-        del q, k, v
         # compute activation in mlp stream, cat again and run second linear layer
-        if self.yak_mlp:
-            mlp = self.mlp_act(mlp[..., self.mlp_hidden_dim_first // 2:]) * mlp[..., :self.mlp_hidden_dim_first // 2]
-        else:
-            mlp = self.mlp_act(mlp)
-        output = self.linear2(torch.cat((attn, mlp), 2))
+        output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
         x += apply_mod(output, mod.gate, None, modulation_dims)
         if x.dtype == torch.float16:
             x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
@@ -350,11 +262,11 @@ class SingleStreamBlock(nn.Module):
 
 
 class LastLayer(nn.Module):
-    def __init__(self, hidden_size: int, patch_size: int, out_channels: int, bias=True, dtype=None, device=None, operations=None):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int, dtype=None, device=None, operations=None):
         super().__init__()
         self.norm_final = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.linear = operations.Linear(hidden_size, patch_size * patch_size * out_channels, bias=bias, dtype=dtype, device=device)
-        self.adaLN_modulation = nn.Sequential(nn.SiLU(), operations.Linear(hidden_size, 2 * hidden_size, bias=bias, dtype=dtype, device=device))
+        self.linear = operations.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True, dtype=dtype, device=device)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), operations.Linear(hidden_size, 2 * hidden_size, bias=True, dtype=dtype, device=device))
 
     def forward(self, x: Tensor, vec: Tensor, modulation_dims=None) -> Tensor:
         if vec.ndim == 2:

comfy/ldm/flux/math.py

@@ -7,8 +7,7 @@ import comfy.model_management
 
 
 def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask=None, transformer_options={}) -> Tensor:
-    if pe is not None:
-        q, k = apply_rope(q, k, pe)
+    q, k = apply_rope(q, k, pe)
     heads = q.shape[1]
     x = optimized_attention(q, k, v, heads, skip_reshape=True, mask=mask, transformer_options=transformer_options)
     return x

comfy/ldm/flux/model.py

@@ -15,8 +15,6 @@ from .layers import (
     MLPEmbedder,
     SingleStreamBlock,
     timestep_embedding,
-    Modulation,
-    RMSNorm
 )
 
 @dataclass
@@ -35,14 +33,6 @@ class FluxParams:
     patch_size: int
     qkv_bias: bool
     guidance_embed: bool
-    txt_ids_dims: list
-    global_modulation: bool = False
-    mlp_silu_act: bool = False
-    ops_bias: bool = True
-    default_ref_method: str = "offset"
-    ref_index_scale: float = 1.0
-    yak_mlp: bool = False
-    txt_norm: bool = False
 
 
 class Flux(nn.Module):
@@ -68,22 +58,13 @@ class Flux(nn.Module):
         self.hidden_size = params.hidden_size
         self.num_heads = params.num_heads
         self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
-        self.img_in = operations.Linear(self.in_channels, self.hidden_size, bias=params.ops_bias, dtype=dtype, device=device)
-        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, bias=params.ops_bias, dtype=dtype, device=device, operations=operations)
-        if params.vec_in_dim is not None:
-            self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size, dtype=dtype, device=device, operations=operations)
-        else:
-            self.vector_in = None
-
+        self.img_in = operations.Linear(self.in_channels, self.hidden_size, bias=True, dtype=dtype, device=device)
+        self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, dtype=dtype, device=device, operations=operations)
+        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size, dtype=dtype, device=device, operations=operations)
         self.guidance_in = (
-            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, bias=params.ops_bias, dtype=dtype, device=device, operations=operations) if params.guidance_embed else nn.Identity()
+            MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, dtype=dtype, device=device, operations=operations) if params.guidance_embed else nn.Identity()
         )
-        self.txt_in = operations.Linear(params.context_in_dim, self.hidden_size, bias=params.ops_bias, dtype=dtype, device=device)
-
-        if params.txt_norm:
-            self.txt_norm = RMSNorm(params.context_in_dim, dtype=dtype, device=device, operations=operations)
-        else:
-            self.txt_norm = None
+        self.txt_in = operations.Linear(params.context_in_dim, self.hidden_size, dtype=dtype, device=device)
 
         self.double_blocks = nn.ModuleList(
             [
@@ -92,10 +73,6 @@ class Flux(nn.Module):
                     self.num_heads,
                     mlp_ratio=params.mlp_ratio,
                     qkv_bias=params.qkv_bias,
-                    modulation=params.global_modulation is False,
-                    mlp_silu_act=params.mlp_silu_act,
-                    proj_bias=params.ops_bias,
-                    yak_mlp=params.yak_mlp,
                     dtype=dtype, device=device, operations=operations
                 )
                 for _ in range(params.depth)
@@ -104,30 +81,13 @@ class Flux(nn.Module):
 
         self.single_blocks = nn.ModuleList(
             [
-                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio, modulation=params.global_modulation is False, mlp_silu_act=params.mlp_silu_act, bias=params.ops_bias, yak_mlp=params.yak_mlp, dtype=dtype, device=device, operations=operations)
+                SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio, dtype=dtype, device=device, operations=operations)
                 for _ in range(params.depth_single_blocks)
             ]
         )
 
         if final_layer:
-            self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, bias=params.ops_bias, dtype=dtype, device=device, operations=operations)
-
-        if params.global_modulation:
-            self.double_stream_modulation_img = Modulation(
-                self.hidden_size,
-                double=True,
-                bias=False,
-                dtype=dtype, device=device, operations=operations
-            )
-            self.double_stream_modulation_txt = Modulation(
-                self.hidden_size,
-                double=True,
-                bias=False,
-                dtype=dtype, device=device, operations=operations
-            )
-            self.single_stream_modulation = Modulation(
-                self.hidden_size, double=False, bias=False, dtype=dtype, device=device, operations=operations
-            )
+            self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations)
 
     def forward_orig(
         self,
@@ -143,6 +103,9 @@ class Flux(nn.Module):
         attn_mask: Tensor = None,
     ) -> Tensor:
 
+        if y is None:
+            y = torch.zeros((img.shape[0], self.params.vec_in_dim), device=img.device, dtype=img.dtype)
+
         patches = transformer_options.get("patches", {})
         patches_replace = transformer_options.get("patches_replace", {})
         if img.ndim != 3 or txt.ndim != 3:
@@ -155,19 +118,9 @@ class Flux(nn.Module):
             if guidance is not None:
                 vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
 
-        if self.vector_in is not None:
-            if y is None:
-                y = torch.zeros((img.shape[0], self.params.vec_in_dim), device=img.device, dtype=img.dtype)
-            vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
-
-        if self.txt_norm is not None:
-            txt = self.txt_norm(txt)
+        vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
         txt = self.txt_in(txt)
 
-        vec_orig = vec
-        if self.params.global_modulation:
-            vec = (self.double_stream_modulation_img(vec_orig), self.double_stream_modulation_txt(vec_orig))
-
         if "post_input" in patches:
             for p in patches["post_input"]:
                 out = p({"img": img, "txt": txt, "img_ids": img_ids, "txt_ids": txt_ids})
@@ -183,10 +136,7 @@ class Flux(nn.Module):
             pe = None
 
         blocks_replace = patches_replace.get("dit", {})
-        transformer_options["total_blocks"] = len(self.double_blocks)
-        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.double_blocks):
-            transformer_options["block_index"] = i
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -227,13 +177,7 @@ class Flux(nn.Module):
 
         img = torch.cat((txt, img), 1)
 
-        if self.params.global_modulation:
-            vec, _ = self.single_stream_modulation(vec_orig)
-
-        transformer_options["total_blocks"] = len(self.single_blocks)
-        transformer_options["block_type"] = "single"
         for i, block in enumerate(self.single_blocks):
-            transformer_options["block_index"] = i
             if ("single_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -263,10 +207,10 @@ class Flux(nn.Module):
 
         img = img[:, txt.shape[1] :, ...]
 
-        img = self.final_layer(img, vec_orig)  # (N, T, patch_size ** 2 * out_channels)
+        img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
         return img
 
-    def process_img(self, x, index=0, h_offset=0, w_offset=0, transformer_options={}):
+    def process_img(self, x, index=0, h_offset=0, w_offset=0):
         bs, c, h, w = x.shape
         patch_size = self.patch_size
         x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
@@ -278,22 +222,10 @@ class Flux(nn.Module):
         h_offset = ((h_offset + (patch_size // 2)) // patch_size)
         w_offset = ((w_offset + (patch_size // 2)) // patch_size)
 
-        steps_h = h_len
-        steps_w = w_len
-
-        rope_options = transformer_options.get("rope_options", None)
-        if rope_options is not None:
-            h_len = (h_len - 1.0) * rope_options.get("scale_y", 1.0) + 1.0
-            w_len = (w_len - 1.0) * rope_options.get("scale_x", 1.0) + 1.0
-
-            index += rope_options.get("shift_t", 0.0)
-            h_offset += rope_options.get("shift_y", 0.0)
-            w_offset += rope_options.get("shift_x", 0.0)
-
-        img_ids = torch.zeros((steps_h, steps_w, len(self.params.axes_dim)), device=x.device, dtype=torch.float32)
+        img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
         img_ids[:, :, 0] = img_ids[:, :, 1] + index
-        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(h_offset, h_len - 1 + h_offset, steps=steps_h, device=x.device, dtype=torch.float32).unsqueeze(1)
-        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(w_offset, w_len - 1 + w_offset, steps=steps_w, device=x.device, dtype=torch.float32).unsqueeze(0)
+        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(h_offset, h_len - 1 + h_offset, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1)
+        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(w_offset, w_len - 1 + w_offset, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0)
         return img, repeat(img_ids, "h w c -> b (h w) c", b=bs)
 
     def forward(self, x, timestep, context, y=None, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs):
@@ -309,16 +241,16 @@ class Flux(nn.Module):
 
         h_len = ((h_orig + (patch_size // 2)) // patch_size)
         w_len = ((w_orig + (patch_size // 2)) // patch_size)
-        img, img_ids = self.process_img(x, transformer_options=transformer_options)
+        img, img_ids = self.process_img(x)
         img_tokens = img.shape[1]
         if ref_latents is not None:
             h = 0
             w = 0
             index = 0
-            ref_latents_method = kwargs.get("ref_latents_method", self.params.default_ref_method)
+            ref_latents_method = kwargs.get("ref_latents_method", "offset")
             for ref in ref_latents:
                 if ref_latents_method == "index":
-                    index += self.params.ref_index_scale
+                    index += 1
                     h_offset = 0
                     w_offset = 0
                 elif ref_latents_method == "uxo":
@@ -342,12 +274,7 @@ class Flux(nn.Module):
                 img = torch.cat([img, kontext], dim=1)
                 img_ids = torch.cat([img_ids, kontext_ids], dim=1)
 
-        txt_ids = torch.zeros((bs, context.shape[1], len(self.params.axes_dim)), device=x.device, dtype=torch.float32)
-
-        if len(self.params.txt_ids_dims) > 0:
-            for i in self.params.txt_ids_dims:
-                txt_ids[:, :, i] = torch.linspace(0, context.shape[1] - 1, steps=context.shape[1], device=x.device, dtype=torch.float32)
-
+        txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
         out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control, transformer_options, attn_mask=kwargs.get("attention_mask", None))
         out = out[:, :img_tokens]
-        return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=self.patch_size, pw=self.patch_size)[:,:,:h_orig,:w_orig]
+        return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h_orig,:w_orig]

comfy/ldm/hunyuan_video/model.py

@@ -6,6 +6,7 @@ import comfy.ldm.flux.layers
 import comfy.ldm.modules.diffusionmodules.mmdit
 from comfy.ldm.modules.attention import optimized_attention
 
+
 from dataclasses import dataclass
 from einops import repeat
 
@@ -41,9 +42,6 @@ class HunyuanVideoParams:
     guidance_embed: bool
     byt5: bool
     meanflow: bool
-    use_cond_type_embedding: bool
-    vision_in_dim: int
-    meanflow_sum: bool
 
 
 class SelfAttentionRef(nn.Module):
@@ -159,10 +157,7 @@ class TokenRefiner(nn.Module):
         t = self.t_embedder(timestep_embedding(timesteps, 256, time_factor=1.0).to(x.dtype))
         # m = mask.float().unsqueeze(-1)
         # c = (x.float() * m).sum(dim=1) / m.sum(dim=1) #TODO: the following works when the x.shape is the same length as the tokens but might break otherwise
-        if x.dtype == torch.float16:
-            c = x.float().sum(dim=1) / x.shape[1]
-        else:
-            c = x.sum(dim=1) / x.shape[1]
+        c = x.sum(dim=1) / x.shape[1]
 
         c = t + self.c_embedder(c.to(x.dtype))
         x = self.input_embedder(x)
@@ -201,15 +196,11 @@ class HunyuanVideo(nn.Module):
     def __init__(self, image_model=None, final_layer=True, dtype=None, device=None, operations=None, **kwargs):
         super().__init__()
         self.dtype = dtype
-        operation_settings = {"operations": operations, "device": device, "dtype": dtype}
-
         params = HunyuanVideoParams(**kwargs)
         self.params = params
         self.patch_size = params.patch_size
         self.in_channels = params.in_channels
         self.out_channels = params.out_channels
-        self.use_cond_type_embedding = params.use_cond_type_embedding
-        self.vision_in_dim = params.vision_in_dim
         if params.hidden_size % params.num_heads != 0:
             raise ValueError(
                 f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
@@ -275,18 +266,6 @@ class HunyuanVideo(nn.Module):
         if final_layer:
             self.final_layer = LastLayer(self.hidden_size, self.patch_size[-1], self.out_channels, dtype=dtype, device=device, operations=operations)
 
-        # HunyuanVideo 1.5 specific modules
-        if self.vision_in_dim is not None:
-            from comfy.ldm.wan.model import MLPProj
-            self.vision_in = MLPProj(in_dim=self.vision_in_dim, out_dim=self.hidden_size, operation_settings=operation_settings)
-        else:
-            self.vision_in = None
-        if self.use_cond_type_embedding:
-            # 0: text_encoder feature 1: byt5 feature 2: vision_encoder feature
-            self.cond_type_embedding = nn.Embedding(3, self.hidden_size)
-        else:
-            self.cond_type_embedding = None
-
     def forward_orig(
         self,
         img: Tensor,
@@ -297,7 +276,6 @@ class HunyuanVideo(nn.Module):
         timesteps: Tensor,
         y: Tensor = None,
         txt_byt5=None,
-        clip_fea=None,
         guidance: Tensor = None,
         guiding_frame_index=None,
         ref_latent=None,
@@ -318,7 +296,7 @@ class HunyuanVideo(nn.Module):
                 timesteps_r = transformer_options['sample_sigmas'][w[0] + 1]
                 timesteps_r = timesteps_r.unsqueeze(0).to(device=timesteps.device, dtype=timesteps.dtype)
                 vec_r = self.time_r_in(timestep_embedding(timesteps_r, 256, time_factor=1000.0).to(img.dtype))
-                vec = (vec + vec_r) if self.params.meanflow_sum else (vec + vec_r) / 2
+                vec = (vec + vec_r) / 2
 
         if ref_latent is not None:
             ref_latent_ids = self.img_ids(ref_latent)
@@ -353,31 +331,12 @@ class HunyuanVideo(nn.Module):
 
         txt = self.txt_in(txt, timesteps, txt_mask, transformer_options=transformer_options)
 
-        if self.cond_type_embedding is not None:
-            self.cond_type_embedding.to(txt.device)
-            cond_emb = self.cond_type_embedding(torch.zeros_like(txt[:, :, 0], device=txt.device, dtype=torch.long))
-            txt = txt + cond_emb.to(txt.dtype)
-
         if self.byt5_in is not None and txt_byt5 is not None:
             txt_byt5 = self.byt5_in(txt_byt5)
-            if self.cond_type_embedding is not None:
-                cond_emb = self.cond_type_embedding(torch.ones_like(txt_byt5[:, :, 0], device=txt_byt5.device, dtype=torch.long))
-                txt_byt5 = txt_byt5 + cond_emb.to(txt_byt5.dtype)
-                txt = torch.cat((txt_byt5, txt), dim=1) # byt5 first for HunyuanVideo1.5
-            else:
-                txt = torch.cat((txt, txt_byt5), dim=1)
             txt_byt5_ids = torch.zeros((txt_ids.shape[0], txt_byt5.shape[1], txt_ids.shape[-1]), device=txt_ids.device, dtype=txt_ids.dtype)
+            txt = torch.cat((txt, txt_byt5), dim=1)
             txt_ids = torch.cat((txt_ids, txt_byt5_ids), dim=1)
 
-        if clip_fea is not None:
-            txt_vision_states = self.vision_in(clip_fea)
-            if self.cond_type_embedding is not None:
-                cond_emb = self.cond_type_embedding(2 * torch.ones_like(txt_vision_states[:, :, 0], dtype=torch.long, device=txt_vision_states.device))
-                txt_vision_states = txt_vision_states + cond_emb
-            txt = torch.cat((txt_vision_states.to(txt.dtype), txt), dim=1)
-            extra_txt_ids = torch.zeros((txt_ids.shape[0], txt_vision_states.shape[1], txt_ids.shape[-1]), device=txt_ids.device, dtype=txt_ids.dtype)
-            txt_ids = torch.cat((txt_ids, extra_txt_ids), dim=1)
-
         ids = torch.cat((img_ids, txt_ids), dim=1)
         pe = self.pe_embedder(ids)
 
@@ -390,10 +349,7 @@ class HunyuanVideo(nn.Module):
             attn_mask = None
 
         blocks_replace = patches_replace.get("dit", {})
-        transformer_options["total_blocks"] = len(self.double_blocks)
-        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.double_blocks):
-            transformer_options["block_index"] = i
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -415,10 +371,7 @@ class HunyuanVideo(nn.Module):
 
         img = torch.cat((img, txt), 1)
 
-        transformer_options["total_blocks"] = len(self.single_blocks)
-        transformer_options["block_type"] = "single"
         for i, block in enumerate(self.single_blocks):
-            transformer_options["block_index"] = i
             if ("single_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -477,14 +430,14 @@ class HunyuanVideo(nn.Module):
         img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0)
         return repeat(img_ids, "h w c -> b (h w) c", b=bs)
 
-    def forward(self, x, timestep, context, y=None, txt_byt5=None, clip_fea=None, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, disable_time_r=False, control=None, transformer_options={}, **kwargs):
+    def forward(self, x, timestep, context, y=None, txt_byt5=None, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, disable_time_r=False, control=None, transformer_options={}, **kwargs):
         return comfy.patcher_extension.WrapperExecutor.new_class_executor(
             self._forward,
             self,
             comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
-        ).execute(x, timestep, context, y, txt_byt5, clip_fea, guidance, attention_mask, guiding_frame_index, ref_latent, disable_time_r, control, transformer_options, **kwargs)
+        ).execute(x, timestep, context, y, txt_byt5, guidance, attention_mask, guiding_frame_index, ref_latent, disable_time_r, control, transformer_options, **kwargs)
 
-    def _forward(self, x, timestep, context, y=None, txt_byt5=None, clip_fea=None, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, disable_time_r=False, control=None, transformer_options={}, **kwargs):
+    def _forward(self, x, timestep, context, y=None, txt_byt5=None, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, disable_time_r=False, control=None, transformer_options={}, **kwargs):
         bs = x.shape[0]
         if len(self.patch_size) == 3:
             img_ids = self.img_ids(x)
@@ -492,5 +445,5 @@ class HunyuanVideo(nn.Module):
         else:
             img_ids = self.img_ids_2d(x)
             txt_ids = torch.zeros((bs, context.shape[1], 2), device=x.device, dtype=x.dtype)
-        out = self.forward_orig(x, img_ids, context, txt_ids, attention_mask, timestep, y, txt_byt5, clip_fea, guidance, guiding_frame_index, ref_latent, disable_time_r=disable_time_r, control=control, transformer_options=transformer_options)
+        out = self.forward_orig(x, img_ids, context, txt_ids, attention_mask, timestep, y, txt_byt5, guidance, guiding_frame_index, ref_latent, disable_time_r=disable_time_r, control=control, transformer_options=transformer_options)
         return out

comfy/ldm/hunyuan_video/upsampler.py

@@ -1,121 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from comfy.ldm.modules.diffusionmodules.model import ResnetBlock, VideoConv3d
-from comfy.ldm.hunyuan_video.vae_refiner import RMS_norm
-import model_management, model_patcher
-
-class SRResidualCausalBlock3D(nn.Module):
-    def __init__(self, channels: int):
-        super().__init__()
-        self.block = nn.Sequential(
-            VideoConv3d(channels, channels, kernel_size=3),
-            nn.SiLU(inplace=True),
-            VideoConv3d(channels, channels, kernel_size=3),
-            nn.SiLU(inplace=True),
-            VideoConv3d(channels, channels, kernel_size=3),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return x + self.block(x)
-
-class SRModel3DV2(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        hidden_channels: int = 64,
-        num_blocks: int = 6,
-        global_residual: bool = False,
-    ):
-        super().__init__()
-        self.in_conv = VideoConv3d(in_channels, hidden_channels, kernel_size=3)
-        self.blocks = nn.ModuleList([SRResidualCausalBlock3D(hidden_channels) for _ in range(num_blocks)])
-        self.out_conv = VideoConv3d(hidden_channels, out_channels, kernel_size=3)
-        self.global_residual = bool(global_residual)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        residual = x
-        y = self.in_conv(x)
-        for blk in self.blocks:
-            y = blk(y)
-        y = self.out_conv(y)
-        if self.global_residual and (y.shape == residual.shape):
-            y = y + residual
-        return y
-
-
-class Upsampler(nn.Module):
-    def __init__(
-        self,
-        z_channels: int,
-        out_channels: int,
-        block_out_channels: tuple[int, ...],
-        num_res_blocks: int = 2,
-    ):
-        super().__init__()
-        self.num_res_blocks = num_res_blocks
-        self.block_out_channels = block_out_channels
-        self.z_channels = z_channels
-
-        ch = block_out_channels[0]
-        self.conv_in = VideoConv3d(z_channels, ch, kernel_size=3)
-
-        self.up = nn.ModuleList()
-
-        for i, tgt in enumerate(block_out_channels):
-            stage = nn.Module()
-            stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
-                                                    out_channels=tgt,
-                                                    temb_channels=0,
-                                                    conv_shortcut=False,
-                                                    conv_op=VideoConv3d, norm_op=RMS_norm)
-                                        for j in range(num_res_blocks + 1)])
-            ch = tgt
-            self.up.append(stage)
-
-        self.norm_out = RMS_norm(ch)
-        self.conv_out = VideoConv3d(ch, out_channels, kernel_size=3)
-
-    def forward(self, z):
-        """
-        Args:
-            z: (B, C, T, H, W)
-            target_shape: (H, W)
-        """
-        # z to block_in
-        repeats = self.block_out_channels[0] // (self.z_channels)
-        x = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1)
-
-        # upsampling
-        for stage in self.up:
-            for blk in stage.block:
-                x = blk(x)
-
-        out = self.conv_out(F.silu(self.norm_out(x)))
-        return out
-
-UPSAMPLERS = {
-    "720p": SRModel3DV2,
-    "1080p": Upsampler,
-}
-
-class HunyuanVideo15SRModel():
-    def __init__(self, model_type, config):
-        self.load_device = model_management.vae_device()
-        offload_device = model_management.vae_offload_device()
-        self.dtype = model_management.vae_dtype(self.load_device)
-        self.model_class = UPSAMPLERS.get(model_type)
-        self.model = self.model_class(**config).eval()
-
-        self.patcher = model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
-
-    def load_sd(self, sd):
-        return self.model.load_state_dict(sd, strict=True)
-
-    def get_sd(self):
-        return self.model.state_dict()
-
-    def resample_latent(self, latent):
-        model_management.load_model_gpu(self.patcher)
-        return self.model(latent.to(self.load_device))

comfy/ldm/hunyuan_video/vae_refiner.py

@@ -1,13 +1,11 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from comfy.ldm.modules.diffusionmodules.model import ResnetBlock, AttnBlock, CarriedConv3d, Normalize, conv_carry_causal_3d, torch_cat_if_needed
+from comfy.ldm.modules.diffusionmodules.model import ResnetBlock, AttnBlock, VideoConv3d, Normalize
 import comfy.ops
 import comfy.ldm.models.autoencoder
-import comfy.model_management
 ops = comfy.ops.disable_weight_init
 
-
 class RMS_norm(nn.Module):
     def __init__(self, dim):
         super().__init__()
@@ -16,10 +14,10 @@ class RMS_norm(nn.Module):
         self.gamma = nn.Parameter(torch.empty(shape))
 
     def forward(self, x):
-        return F.normalize(x, dim=1) * self.scale * comfy.model_management.cast_to(self.gamma, dtype=x.dtype, device=x.device)
+        return F.normalize(x, dim=1) * self.scale * self.gamma
 
 class DnSmpl(nn.Module):
-    def __init__(self, ic, oc, tds, refiner_vae, op):
+    def __init__(self, ic, oc, tds=True, refiner_vae=True, op=VideoConv3d):
         super().__init__()
         fct = 2 * 2 * 2 if tds else 1 * 2 * 2
         assert oc % fct == 0
@@ -29,12 +27,11 @@ class DnSmpl(nn.Module):
         self.tds = tds
         self.gs = fct * ic // oc
 
-    def forward(self, x, conv_carry_in=None, conv_carry_out=None):
+    def forward(self, x):
         r1 = 2 if self.tds else 1
-        h = conv_carry_causal_3d([x], self.conv, conv_carry_in, conv_carry_out)
-
-        if self.tds and self.refiner_vae and conv_carry_in is None:
+        h = self.conv(x)
 
+        if self.tds and self.refiner_vae:
             hf = h[:, :, :1, :, :]
             b, c, f, ht, wd = hf.shape
             hf = hf.reshape(b, c, f, ht // 2, 2, wd // 2, 2)
@@ -42,7 +39,14 @@ class DnSmpl(nn.Module):
             hf = hf.reshape(b, 2 * 2 * c, f, ht // 2, wd // 2)
             hf = torch.cat([hf, hf], dim=1)
 
-            h = h[:, :, 1:, :, :]
+            hn = h[:, :, 1:, :, :]
+            b, c, frms, ht, wd = hn.shape
+            nf = frms // r1
+            hn = hn.reshape(b, c, nf, r1, ht // 2, 2, wd // 2, 2)
+            hn = hn.permute(0, 3, 5, 7, 1, 2, 4, 6)
+            hn = hn.reshape(b, r1 * 2 * 2 * c, nf, ht // 2, wd // 2)
+
+            h = torch.cat([hf, hn], dim=2)
 
             xf = x[:, :, :1, :, :]
             b, ci, f, ht, wd = xf.shape
@@ -50,36 +54,38 @@ class DnSmpl(nn.Module):
             xf = xf.permute(0, 4, 6, 1, 2, 3, 5)
             xf = xf.reshape(b, 2 * 2 * ci, f, ht // 2, wd // 2)
             B, C, T, H, W = xf.shape
-            xf = xf.view(B, hf.shape[1], self.gs // 2, T, H, W).mean(dim=2)
+            xf = xf.view(B, h.shape[1], self.gs // 2, T, H, W).mean(dim=2)
 
-            x = x[:, :, 1:, :, :]
+            xn = x[:, :, 1:, :, :]
+            b, ci, frms, ht, wd = xn.shape
+            nf = frms // r1
+            xn = xn.reshape(b, ci, nf, r1, ht // 2, 2, wd // 2, 2)
+            xn = xn.permute(0, 3, 5, 7, 1, 2, 4, 6)
+            xn = xn.reshape(b, r1 * 2 * 2 * ci, nf, ht // 2, wd // 2)
+            B, C, T, H, W = xn.shape
+            xn = xn.view(B, h.shape[1], self.gs, T, H, W).mean(dim=2)
+            sc = torch.cat([xf, xn], dim=2)
+        else:
+            b, c, frms, ht, wd = h.shape
 
-        if h.shape[2] == 0:
-            return hf + xf
+            nf = frms // r1
+            h = h.reshape(b, c, nf, r1, ht // 2, 2, wd // 2, 2)
+            h = h.permute(0, 3, 5, 7, 1, 2, 4, 6)
+            h = h.reshape(b, r1 * 2 * 2 * c, nf, ht // 2, wd // 2)
 
-        b, c, frms, ht, wd = h.shape
-        nf = frms // r1
-        h = h.reshape(b, c, nf, r1, ht // 2, 2, wd // 2, 2)
-        h = h.permute(0, 3, 5, 7, 1, 2, 4, 6)
-        h = h.reshape(b, r1 * 2 * 2 * c, nf, ht // 2, wd // 2)
+            b, ci, frms, ht, wd = x.shape
+            nf = frms // r1
+            sc = x.reshape(b, ci, nf, r1, ht // 2, 2, wd // 2, 2)
+            sc = sc.permute(0, 3, 5, 7, 1, 2, 4, 6)
+            sc = sc.reshape(b, r1 * 2 * 2 * ci, nf, ht // 2, wd // 2)
+            B, C, T, H, W = sc.shape
+            sc = sc.view(B, h.shape[1], self.gs, T, H, W).mean(dim=2)
 
-        b, ci, frms, ht, wd = x.shape
-        nf = frms // r1
-        x = x.reshape(b, ci, nf, r1, ht // 2, 2, wd // 2, 2)
-        x = x.permute(0, 3, 5, 7, 1, 2, 4, 6)
-        x = x.reshape(b, r1 * 2 * 2 * ci, nf, ht // 2, wd // 2)
-        B, C, T, H, W = x.shape
-        x = x.view(B, h.shape[1], self.gs, T, H, W).mean(dim=2)
-
-        if self.tds and self.refiner_vae and conv_carry_in is None:
-            h = torch.cat([hf, h], dim=2)
-            x = torch.cat([xf, x], dim=2)
-
-        return h + x
+        return h + sc
 
 
 class UpSmpl(nn.Module):
-    def __init__(self, ic, oc, tus, refiner_vae, op):
+    def __init__(self, ic, oc, tus=True, refiner_vae=True, op=VideoConv3d):
         super().__init__()
         fct = 2 * 2 * 2 if tus else 1 * 2 * 2
         self.conv = op(ic, oc * fct, kernel_size=3, stride=1, padding=1)
@@ -88,11 +94,11 @@ class UpSmpl(nn.Module):
         self.tus = tus
         self.rp = fct * oc // ic
 
-    def forward(self, x, conv_carry_in=None, conv_carry_out=None):
+    def forward(self, x):
         r1 = 2 if self.tus else 1
-        h = conv_carry_causal_3d([x], self.conv, conv_carry_in, conv_carry_out)
+        h = self.conv(x)
 
-        if self.tus and self.refiner_vae and conv_carry_in is None:
+        if self.tus and self.refiner_vae:
             hf = h[:, :, :1, :, :]
             b, c, f, ht, wd = hf.shape
             nc = c // (2 * 2)
@@ -101,7 +107,14 @@ class UpSmpl(nn.Module):
             hf = hf.reshape(b, nc, f, ht * 2, wd * 2)
             hf = hf[:, : hf.shape[1] // 2]
 
-            h = h[:, :, 1:, :, :]
+            hn = h[:, :, 1:, :, :]
+            b, c, frms, ht, wd = hn.shape
+            nc = c // (r1 * 2 * 2)
+            hn = hn.reshape(b, r1, 2, 2, nc, frms, ht, wd)
+            hn = hn.permute(0, 4, 5, 1, 6, 2, 7, 3)
+            hn = hn.reshape(b, nc, frms * r1, ht * 2, wd * 2)
+
+            h = torch.cat([hf, hn], dim=2)
 
             xf = x[:, :, :1, :, :]
             b, ci, f, ht, wd = xf.shape
@@ -112,26 +125,29 @@ class UpSmpl(nn.Module):
             xf = xf.permute(0, 3, 4, 5, 1, 6, 2)
             xf = xf.reshape(b, nc, f, ht * 2, wd * 2)
 
-            x = x[:, :, 1:, :, :]
+            xn = x[:, :, 1:, :, :]
+            xn = xn.repeat_interleave(repeats=self.rp, dim=1)
+            b, c, frms, ht, wd = xn.shape
+            nc = c // (r1 * 2 * 2)
+            xn = xn.reshape(b, r1, 2, 2, nc, frms, ht, wd)
+            xn = xn.permute(0, 4, 5, 1, 6, 2, 7, 3)
+            xn = xn.reshape(b, nc, frms * r1, ht * 2, wd * 2)
+            sc = torch.cat([xf, xn], dim=2)
+        else:
+            b, c, frms, ht, wd = h.shape
+            nc = c // (r1 * 2 * 2)
+            h = h.reshape(b, r1, 2, 2, nc, frms, ht, wd)
+            h = h.permute(0, 4, 5, 1, 6, 2, 7, 3)
+            h = h.reshape(b, nc, frms * r1, ht * 2, wd * 2)
 
-        b, c, frms, ht, wd = h.shape
-        nc = c // (r1 * 2 * 2)
-        h = h.reshape(b, r1, 2, 2, nc, frms, ht, wd)
-        h = h.permute(0, 4, 5, 1, 6, 2, 7, 3)
-        h = h.reshape(b, nc, frms * r1, ht * 2, wd * 2)
+            sc = x.repeat_interleave(repeats=self.rp, dim=1)
+            b, c, frms, ht, wd = sc.shape
+            nc = c // (r1 * 2 * 2)
+            sc = sc.reshape(b, r1, 2, 2, nc, frms, ht, wd)
+            sc = sc.permute(0, 4, 5, 1, 6, 2, 7, 3)
+            sc = sc.reshape(b, nc, frms * r1, ht * 2, wd * 2)
 
-        x = x.repeat_interleave(repeats=self.rp, dim=1)
-        b, c, frms, ht, wd = x.shape
-        nc = c // (r1 * 2 * 2)
-        x = x.reshape(b, r1, 2, 2, nc, frms, ht, wd)
-        x = x.permute(0, 4, 5, 1, 6, 2, 7, 3)
-        x = x.reshape(b, nc, frms * r1, ht * 2, wd * 2)
-
-        if self.tus and self.refiner_vae and conv_carry_in is None:
-            h = torch.cat([hf, h], dim=2)
-            x = torch.cat([xf, x], dim=2)
-
-        return h + x
+        return h + sc
 
 class Encoder(nn.Module):
     def __init__(self, in_channels, z_channels, block_out_channels, num_res_blocks,
@@ -144,7 +160,7 @@ class Encoder(nn.Module):
 
         self.refiner_vae = refiner_vae
         if self.refiner_vae:
-            conv_op = CarriedConv3d
+            conv_op = VideoConv3d
             norm_op = RMS_norm
         else:
             conv_op = ops.Conv3d
@@ -172,9 +188,9 @@ class Encoder(nn.Module):
             self.down.append(stage)
 
         self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, conv_op=conv_op, norm_op=norm_op)
+        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
         self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv3d, norm_op=norm_op)
-        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, conv_op=conv_op, norm_op=norm_op)
+        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
 
         self.norm_out = norm_op(ch)
         self.conv_out = conv_op(ch, z_channels << 1, 3, 1, 1)
@@ -185,48 +201,31 @@ class Encoder(nn.Module):
         if not self.refiner_vae and x.shape[2] == 1:
             x = x.expand(-1, -1, self.ffactor_temporal, -1, -1)
 
-        if self.refiner_vae:
-            xl = [x[:, :, :1, :, :]]
-            if x.shape[2] > self.ffactor_temporal:
-                xl += torch.split(x[:, :, 1: 1 + ((x.shape[2] - 1) // self.ffactor_temporal) * self.ffactor_temporal, :, :], self.ffactor_temporal * 2, dim=2)
-            x = xl
-        else:
-            x = [x]
-        out = []
+        x = self.conv_in(x)
 
-        conv_carry_in = None
+        for stage in self.down:
+            for blk in stage.block:
+                x = blk(x)
+            if hasattr(stage, 'downsample'):
+                x = stage.downsample(x)
 
-        for i, x1 in enumerate(x):
-            conv_carry_out = []
-            if i == len(x) - 1:
-                conv_carry_out = None
-
-            x1 = [ x1 ]
-            x1 = conv_carry_causal_3d(x1, self.conv_in, conv_carry_in, conv_carry_out)
-
-            for stage in self.down:
-                for blk in stage.block:
-                    x1 = blk(x1, None, conv_carry_in, conv_carry_out)
-                if hasattr(stage, 'downsample'):
-                    x1 = stage.downsample(x1, conv_carry_in, conv_carry_out)
-
-            out.append(x1)
-            conv_carry_in = conv_carry_out
-
-        out = torch_cat_if_needed(out, dim=2)
-
-        x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(out)))
-        del out
+        x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
 
         b, c, t, h, w = x.shape
         grp = c // (self.z_channels << 1)
         skip = x.view(b, c // grp, grp, t, h, w).mean(2)
 
-        out = conv_carry_causal_3d([F.silu(self.norm_out(x))], self.conv_out) + skip
+        out = self.conv_out(F.silu(self.norm_out(x))) + skip
 
         if self.refiner_vae:
             out = self.regul(out)[0]
 
+            out = torch.cat((out[:, :, :1], out), dim=2)
+            out = out.permute(0, 2, 1, 3, 4)
+            b, f_times_2, c, h, w = out.shape
+            out = out.reshape(b, f_times_2 // 2, 2 * c, h, w)
+            out = out.permute(0, 2, 1, 3, 4).contiguous()
+
         return out
 
 class Decoder(nn.Module):
@@ -240,7 +239,7 @@ class Decoder(nn.Module):
 
         self.refiner_vae = refiner_vae
         if self.refiner_vae:
-            conv_op = CarriedConv3d
+            conv_op = VideoConv3d
             norm_op = RMS_norm
         else:
             conv_op = ops.Conv3d
@@ -250,9 +249,9 @@ class Decoder(nn.Module):
         self.conv_in = conv_op(z_channels, ch, kernel_size=3, stride=1, padding=1)
 
         self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, conv_op=conv_op, norm_op=norm_op)
+        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
         self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv3d, norm_op=norm_op)
-        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch,  conv_op=conv_op, norm_op=norm_op)
+        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
 
         self.up = nn.ModuleList()
         depth = (ffactor_spatial >> 1).bit_length()
@@ -276,38 +275,27 @@ class Decoder(nn.Module):
         self.conv_out = conv_op(ch, out_channels, 3, stride=1, padding=1)
 
     def forward(self, z):
-        x = conv_carry_causal_3d([z], self.conv_in) + z.repeat_interleave(self.block_out_channels[0] // self.z_channels, 1)
+        if self.refiner_vae:
+            z = z.permute(0, 2, 1, 3, 4)
+            b, f, c, h, w = z.shape
+            z = z.reshape(b, f, 2, c // 2, h, w)
+            z = z.permute(0, 1, 2, 3, 4, 5).reshape(b, f * 2, c // 2, h, w)
+            z = z.permute(0, 2, 1, 3, 4)
+            z = z[:, :, 1:]
+
+        x = self.conv_in(z) + z.repeat_interleave(self.block_out_channels[0] // self.z_channels, 1)
         x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
 
-        if self.refiner_vae:
-            x = torch.split(x, 2, dim=2)
-        else:
-            x = [ x ]
-        out = []
+        for stage in self.up:
+            for blk in stage.block:
+                x = blk(x)
+            if hasattr(stage, 'upsample'):
+                x = stage.upsample(x)
 
-        conv_carry_in = None
-
-        for i, x1 in enumerate(x):
-            conv_carry_out = []
-            if i == len(x) - 1:
-                conv_carry_out = None
-            for stage in self.up:
-                for blk in stage.block:
-                    x1 = blk(x1, None, conv_carry_in, conv_carry_out)
-                if hasattr(stage, 'upsample'):
-                    x1 = stage.upsample(x1, conv_carry_in, conv_carry_out)
-
-            x1 = [ F.silu(self.norm_out(x1)) ]
-            x1 = conv_carry_causal_3d(x1, self.conv_out, conv_carry_in, conv_carry_out)
-            out.append(x1)
-            conv_carry_in = conv_carry_out
-        del x
-
-        out = torch_cat_if_needed(out, dim=2)
+        out = self.conv_out(F.silu(self.norm_out(x)))
 
         if not self.refiner_vae:
             if z.shape[-3] == 1:
                 out = out[:, :, -1:]
 
         return out
-

comfy/ldm/kandinsky5/model.py

@@ -1,413 +0,0 @@
-import torch
-from torch import nn
-import math
-
-import comfy.ldm.common_dit
-from comfy.ldm.modules.attention import optimized_attention
-from comfy.ldm.flux.math import apply_rope1
-from comfy.ldm.flux.layers import EmbedND
-
-def attention(q, k, v, heads, transformer_options={}):
-    return optimized_attention(
-        q.transpose(1, 2),
-        k.transpose(1, 2),
-        v.transpose(1, 2),
-        heads=heads,
-        skip_reshape=True,
-        transformer_options=transformer_options
-    )
-
-def apply_scale_shift_norm(norm, x, scale, shift):
-    return torch.addcmul(shift, norm(x), scale + 1.0)
-
-def apply_gate_sum(x, out, gate):
-    return torch.addcmul(x, gate, out)
-
-def get_shift_scale_gate(params):
-    shift, scale, gate = torch.chunk(params, 3, dim=-1)
-    return tuple(x.unsqueeze(1) for x in (shift, scale, gate))
-
-def get_freqs(dim, max_period=10000.0):
-    return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim)
-
-
-class TimeEmbeddings(nn.Module):
-    def __init__(self, model_dim, time_dim, max_period=10000.0, operation_settings=None):
-        super().__init__()
-        assert model_dim % 2 == 0
-        self.model_dim = model_dim
-        self.max_period = max_period
-        self.register_buffer("freqs", get_freqs(model_dim // 2, max_period), persistent=False)
-        operations = operation_settings.get("operations")
-        self.in_layer = operations.Linear(model_dim, time_dim, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.activation = nn.SiLU()
-        self.out_layer = operations.Linear(time_dim, time_dim, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-
-    def forward(self, timestep, dtype):
-        args = torch.outer(timestep, self.freqs.to(device=timestep.device))
-        time_embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(dtype)
-        time_embed = self.out_layer(self.activation(self.in_layer(time_embed)))
-        return time_embed
-
-
-class TextEmbeddings(nn.Module):
-    def __init__(self, text_dim, model_dim, operation_settings=None):
-        super().__init__()
-        operations = operation_settings.get("operations")
-        self.in_layer = operations.Linear(text_dim, model_dim, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.norm = operations.LayerNorm(model_dim, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-
-    def forward(self, text_embed):
-        text_embed = self.in_layer(text_embed)
-        return self.norm(text_embed).type_as(text_embed)
-
-
-class VisualEmbeddings(nn.Module):
-    def __init__(self, visual_dim, model_dim, patch_size, operation_settings=None):
-        super().__init__()
-        self.patch_size = patch_size
-        operations = operation_settings.get("operations")
-        self.in_layer = operations.Linear(visual_dim, model_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-
-    def forward(self, x):
-        x = x.movedim(1, -1)  # B C T H W -> B T H W C
-        B, T, H, W, dim = x.shape
-        pt, ph, pw = self.patch_size
-
-        x = x.view(
-            B,
-            T // pt, pt,
-            H // ph, ph,
-            W // pw, pw,
-            dim,
-        ).permute(0, 1, 3, 5, 2, 4, 6, 7).flatten(4, 7)
-
-        return self.in_layer(x)
-
-
-class Modulation(nn.Module):
-    def __init__(self, time_dim, model_dim, num_params, operation_settings=None):
-        super().__init__()
-        self.activation = nn.SiLU()
-        self.out_layer = operation_settings.get("operations").Linear(time_dim, num_params * model_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-
-    def forward(self, x):
-        return self.out_layer(self.activation(x))
-
-
-class SelfAttention(nn.Module):
-    def __init__(self, num_channels, head_dim, operation_settings=None):
-        super().__init__()
-        assert num_channels % head_dim == 0
-        self.num_heads = num_channels // head_dim
-        self.head_dim = head_dim
-
-        operations = operation_settings.get("operations")
-        self.to_query = operations.Linear(num_channels, num_channels, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.to_key = operations.Linear(num_channels, num_channels, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.to_value = operations.Linear(num_channels, num_channels, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.query_norm = operations.RMSNorm(head_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.key_norm = operations.RMSNorm(head_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-
-        self.out_layer = operations.Linear(num_channels, num_channels, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.num_chunks = 2
-
-    def _compute_qk(self, x, freqs, proj_fn, norm_fn):
-        result = proj_fn(x).view(*x.shape[:-1], self.num_heads, -1)
-        return apply_rope1(norm_fn(result), freqs)
-
-    def _forward(self, x, freqs, transformer_options={}):
-        q = self._compute_qk(x, freqs, self.to_query, self.query_norm)
-        k = self._compute_qk(x, freqs, self.to_key, self.key_norm)
-        v = self.to_value(x).view(*x.shape[:-1], self.num_heads, -1)
-        out = attention(q, k, v, self.num_heads, transformer_options=transformer_options)
-        return self.out_layer(out)
-
-    def _forward_chunked(self, x, freqs, transformer_options={}):
-        def process_chunks(proj_fn, norm_fn):
-            x_chunks = torch.chunk(x, self.num_chunks, dim=1)
-            freqs_chunks = torch.chunk(freqs, self.num_chunks, dim=1)
-            chunks = []
-            for x_chunk, freqs_chunk in zip(x_chunks, freqs_chunks):
-                chunks.append(self._compute_qk(x_chunk, freqs_chunk, proj_fn, norm_fn))
-            return torch.cat(chunks, dim=1)
-
-        q = process_chunks(self.to_query, self.query_norm)
-        k = process_chunks(self.to_key, self.key_norm)
-        v = self.to_value(x).view(*x.shape[:-1], self.num_heads, -1)
-        out = attention(q, k, v, self.num_heads, transformer_options=transformer_options)
-        return self.out_layer(out)
-
-    def forward(self, x, freqs, transformer_options={}):
-        if x.shape[1] > 8192:
-            return self._forward_chunked(x, freqs, transformer_options=transformer_options)
-        else:
-            return self._forward(x, freqs, transformer_options=transformer_options)
-
-
-class CrossAttention(SelfAttention):
-    def get_qkv(self, x, context):
-        q = self.to_query(x).view(*x.shape[:-1], self.num_heads, -1)
-        k = self.to_key(context).view(*context.shape[:-1], self.num_heads, -1)
-        v = self.to_value(context).view(*context.shape[:-1], self.num_heads, -1)
-        return q, k, v
-
-    def forward(self, x, context, transformer_options={}):
-        q, k, v = self.get_qkv(x, context)
-        out = attention(self.query_norm(q), self.key_norm(k), v, self.num_heads, transformer_options=transformer_options)
-        return self.out_layer(out)
-
-
-class FeedForward(nn.Module):
-    def __init__(self, dim, ff_dim, operation_settings=None):
-        super().__init__()
-        operations = operation_settings.get("operations")
-        self.in_layer = operations.Linear(dim, ff_dim, bias=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.activation = nn.GELU()
-        self.out_layer = operations.Linear(ff_dim, dim, bias=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.num_chunks = 4
-
-    def _forward(self, x):
-        return self.out_layer(self.activation(self.in_layer(x)))
-
-    def _forward_chunked(self, x):
-        chunks = torch.chunk(x, self.num_chunks, dim=1)
-        output_chunks = []
-        for chunk in chunks:
-            output_chunks.append(self._forward(chunk))
-        return torch.cat(output_chunks, dim=1)
-
-    def forward(self, x):
-        if x.shape[1] > 8192:
-            return self._forward_chunked(x)
-        else:
-            return self._forward(x)
-
-
-class OutLayer(nn.Module):
-    def __init__(self, model_dim, time_dim, visual_dim, patch_size, operation_settings=None):
-        super().__init__()
-        self.patch_size = patch_size
-        self.modulation = Modulation(time_dim, model_dim, 2, operation_settings=operation_settings)
-        operations = operation_settings.get("operations")
-        self.norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.out_layer = operations.Linear(model_dim, math.prod(patch_size) * visual_dim, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-
-    def forward(self, visual_embed, time_embed):
-        B, T, H, W, _ = visual_embed.shape
-        shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1)
-        scale = scale[:, None, None, None, :]
-        shift = shift[:, None, None, None, :]
-        visual_embed = apply_scale_shift_norm(self.norm, visual_embed, scale, shift)
-        x = self.out_layer(visual_embed)
-
-        out_dim = x.shape[-1] // (self.patch_size[0] * self.patch_size[1] * self.patch_size[2])
-        x = x.view(
-            B, T, H, W,
-            out_dim,
-            self.patch_size[0], self.patch_size[1], self.patch_size[2]
-        )
-        return x.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(2, 3).flatten(3, 4).flatten(4, 5)
-
-
-class TransformerEncoderBlock(nn.Module):
-    def __init__(self, model_dim, time_dim, ff_dim, head_dim, operation_settings=None):
-        super().__init__()
-        self.text_modulation = Modulation(time_dim, model_dim, 6, operation_settings=operation_settings)
-        operations = operation_settings.get("operations")
-
-        self.self_attention_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.self_attention = SelfAttention(model_dim, head_dim, operation_settings=operation_settings)
-
-        self.feed_forward_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.feed_forward = FeedForward(model_dim, ff_dim, operation_settings=operation_settings)
-
-    def forward(self, x, time_embed, freqs, transformer_options={}):
-        self_attn_params, ff_params = torch.chunk(self.text_modulation(time_embed), 2, dim=-1)
-        shift, scale, gate = get_shift_scale_gate(self_attn_params)
-        out = apply_scale_shift_norm(self.self_attention_norm, x, scale, shift)
-        out = self.self_attention(out, freqs, transformer_options=transformer_options)
-        x = apply_gate_sum(x, out, gate)
-
-        shift, scale, gate = get_shift_scale_gate(ff_params)
-        out = apply_scale_shift_norm(self.feed_forward_norm, x, scale, shift)
-        out = self.feed_forward(out)
-        x = apply_gate_sum(x, out, gate)
-        return x
-
-
-class TransformerDecoderBlock(nn.Module):
-    def __init__(self, model_dim, time_dim, ff_dim, head_dim, operation_settings=None):
-        super().__init__()
-        self.visual_modulation = Modulation(time_dim, model_dim, 9, operation_settings=operation_settings)
-
-        operations = operation_settings.get("operations")
-        self.self_attention_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.self_attention = SelfAttention(model_dim, head_dim, operation_settings=operation_settings)
-
-        self.cross_attention_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.cross_attention = CrossAttention(model_dim, head_dim, operation_settings=operation_settings)
-
-        self.feed_forward_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.feed_forward = FeedForward(model_dim, ff_dim, operation_settings=operation_settings)
-
-    def forward(self, visual_embed, text_embed, time_embed, freqs, transformer_options={}):
-        self_attn_params, cross_attn_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 3, dim=-1)
-        # self attention
-        shift, scale, gate = get_shift_scale_gate(self_attn_params)
-        visual_out = apply_scale_shift_norm(self.self_attention_norm, visual_embed, scale, shift)
-        visual_out = self.self_attention(visual_out, freqs, transformer_options=transformer_options)
-        visual_embed = apply_gate_sum(visual_embed, visual_out, gate)
-        # cross attention
-        shift, scale, gate = get_shift_scale_gate(cross_attn_params)
-        visual_out = apply_scale_shift_norm(self.cross_attention_norm, visual_embed, scale, shift)
-        visual_out = self.cross_attention(visual_out, text_embed, transformer_options=transformer_options)
-        visual_embed = apply_gate_sum(visual_embed, visual_out, gate)
-        # feed forward
-        shift, scale, gate = get_shift_scale_gate(ff_params)
-        visual_out = apply_scale_shift_norm(self.feed_forward_norm, visual_embed, scale, shift)
-        visual_out = self.feed_forward(visual_out)
-        visual_embed = apply_gate_sum(visual_embed, visual_out, gate)
-        return visual_embed
-
-
-class Kandinsky5(nn.Module):
-    def __init__(
-        self,
-        in_visual_dim=16, out_visual_dim=16, in_text_dim=3584, in_text_dim2=768, time_dim=512,
-        model_dim=1792, ff_dim=7168, visual_embed_dim=132, patch_size=(1, 2, 2), num_text_blocks=2, num_visual_blocks=32,
-        axes_dims=(16, 24, 24), rope_scale_factor=(1.0, 2.0, 2.0),
-        dtype=None, device=None, operations=None, **kwargs
-    ):
-        super().__init__()
-        head_dim = sum(axes_dims)
-        self.rope_scale_factor = rope_scale_factor
-        self.in_visual_dim = in_visual_dim
-        self.model_dim = model_dim
-        self.patch_size = patch_size
-        self.visual_embed_dim = visual_embed_dim
-        self.dtype = dtype
-        self.device = device
-        operation_settings = {"operations": operations, "device": device, "dtype": dtype}
-
-        self.time_embeddings = TimeEmbeddings(model_dim, time_dim, operation_settings=operation_settings)
-        self.text_embeddings = TextEmbeddings(in_text_dim, model_dim, operation_settings=operation_settings)
-        self.pooled_text_embeddings = TextEmbeddings(in_text_dim2, time_dim, operation_settings=operation_settings)
-        self.visual_embeddings = VisualEmbeddings(visual_embed_dim, model_dim, patch_size, operation_settings=operation_settings)
-
-        self.text_transformer_blocks = nn.ModuleList(
-            [TransformerEncoderBlock(model_dim, time_dim, ff_dim, head_dim, operation_settings=operation_settings) for _ in range(num_text_blocks)]
-        )
-
-        self.visual_transformer_blocks = nn.ModuleList(
-            [TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim, operation_settings=operation_settings) for _ in range(num_visual_blocks)]
-        )
-
-        self.out_layer = OutLayer(model_dim, time_dim, out_visual_dim, patch_size, operation_settings=operation_settings)
-
-        self.rope_embedder_3d = EmbedND(dim=head_dim, theta=10000.0, axes_dim=axes_dims)
-        self.rope_embedder_1d = EmbedND(dim=head_dim, theta=10000.0, axes_dim=[head_dim])
-
-    def rope_encode_1d(self, seq_len, seq_start=0, steps=None, device=None, dtype=None, transformer_options={}):
-        steps = seq_len if steps is None else steps
-        seq_ids = torch.linspace(seq_start, seq_start + (seq_len - 1), steps=steps, device=device, dtype=dtype)
-        seq_ids = seq_ids.reshape(-1, 1).unsqueeze(0)  # Shape: (1, steps, 1)
-        freqs = self.rope_embedder_1d(seq_ids).movedim(1, 2)
-        return freqs
-
-    def rope_encode_3d(self, t, h, w, t_start=0, steps_t=None, steps_h=None, steps_w=None, device=None, dtype=None, transformer_options={}):
-
-        patch_size = self.patch_size
-        t_len = ((t + (patch_size[0] // 2)) // patch_size[0])
-        h_len = ((h + (patch_size[1] // 2)) // patch_size[1])
-        w_len = ((w + (patch_size[2] // 2)) // patch_size[2])
-
-        if steps_t is None:
-            steps_t = t_len
-        if steps_h is None:
-            steps_h = h_len
-        if steps_w is None:
-            steps_w = w_len
-
-        h_start = 0
-        w_start = 0
-        rope_options = transformer_options.get("rope_options", None)
-        if rope_options is not None:
-            t_len = (t_len - 1.0) * rope_options.get("scale_t", 1.0) + 1.0
-            h_len = (h_len - 1.0) * rope_options.get("scale_y", 1.0) + 1.0
-            w_len = (w_len - 1.0) * rope_options.get("scale_x", 1.0) + 1.0
-
-            t_start += rope_options.get("shift_t", 0.0)
-            h_start += rope_options.get("shift_y", 0.0)
-            w_start += rope_options.get("shift_x", 0.0)
-        else:
-            rope_scale_factor = self.rope_scale_factor
-            if self.model_dim == 4096: # pro video model uses different rope scaling at higher resolutions
-                if h * w >= 14080:
-                    rope_scale_factor = (1.0, 3.16, 3.16)
-
-            t_len = (t_len - 1.0) / rope_scale_factor[0] + 1.0
-            h_len = (h_len - 1.0) / rope_scale_factor[1] + 1.0
-            w_len = (w_len - 1.0) / rope_scale_factor[2] + 1.0
-
-        img_ids = torch.zeros((steps_t, steps_h, steps_w, 3), device=device, dtype=dtype)
-        img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.linspace(t_start, t_start + (t_len - 1), steps=steps_t, device=device, dtype=dtype).reshape(-1, 1, 1)
-        img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.linspace(h_start, h_start + (h_len - 1), steps=steps_h, device=device, dtype=dtype).reshape(1, -1, 1)
-        img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.linspace(w_start, w_start + (w_len - 1), steps=steps_w, device=device, dtype=dtype).reshape(1, 1, -1)
-        img_ids = img_ids.reshape(1, -1, img_ids.shape[-1])
-
-        freqs = self.rope_embedder_3d(img_ids).movedim(1, 2)
-        return freqs
-
-    def forward_orig(self, x, timestep, context, y, freqs, freqs_text, transformer_options={}, **kwargs):
-        patches_replace = transformer_options.get("patches_replace", {})
-        context = self.text_embeddings(context)
-        time_embed = self.time_embeddings(timestep, x.dtype) + self.pooled_text_embeddings(y)
-
-        for block in self.text_transformer_blocks:
-            context = block(context, time_embed, freqs_text, transformer_options=transformer_options)
-
-        visual_embed = self.visual_embeddings(x)
-        visual_shape = visual_embed.shape[:-1]
-        visual_embed = visual_embed.flatten(1, -2)
-
-        blocks_replace = patches_replace.get("dit", {})
-        transformer_options["total_blocks"] = len(self.visual_transformer_blocks)
-        transformer_options["block_type"] = "double"
-        for i, block in enumerate(self.visual_transformer_blocks):
-            transformer_options["block_index"] = i
-            if ("double_block", i) in blocks_replace:
-                def block_wrap(args):
-                    return block(x=args["x"], context=args["context"], time_embed=args["time_embed"], freqs=args["freqs"], transformer_options=args.get("transformer_options"))
-                visual_embed = blocks_replace[("double_block", i)]({"x": visual_embed, "context": context, "time_embed": time_embed, "freqs": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})["x"]
-            else:
-                visual_embed = block(visual_embed, context, time_embed, freqs=freqs, transformer_options=transformer_options)
-
-        visual_embed = visual_embed.reshape(*visual_shape, -1)
-        return self.out_layer(visual_embed, time_embed)
-
-    def _forward(self, x, timestep, context, y, time_dim_replace=None, transformer_options={}, **kwargs):
-        original_dims = x.ndim
-        if original_dims == 4:
-            x = x.unsqueeze(2)
-        bs, c, t_len, h, w = x.shape
-        x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size)
-
-        if time_dim_replace is not None:
-            time_dim_replace = comfy.ldm.common_dit.pad_to_patch_size(time_dim_replace, self.patch_size)
-            x[:, :time_dim_replace.shape[1], :time_dim_replace.shape[2]] = time_dim_replace
-
-        freqs = self.rope_encode_3d(t_len, h, w, device=x.device, dtype=x.dtype, transformer_options=transformer_options)
-        freqs_text = self.rope_encode_1d(context.shape[1], device=x.device, dtype=x.dtype, transformer_options=transformer_options)
-
-        out = self.forward_orig(x, timestep, context, y, freqs, freqs_text, transformer_options=transformer_options, **kwargs)
-        if original_dims == 4:
-            out = out.squeeze(2)
-        return out
-
-    def forward(self, x, timestep, context, y, time_dim_replace=None, transformer_options={}, **kwargs):
-        return comfy.patcher_extension.WrapperExecutor.new_class_executor(
-            self._forward,
-            self,
-            comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
-        ).execute(x, timestep, context, y, time_dim_replace=time_dim_replace, transformer_options=transformer_options, **kwargs)

comfy/ldm/lumina/controlnet.py

@@ -1,160 +0,0 @@
-import torch
-from torch import nn
-
-from .model import JointTransformerBlock
-
-class ZImageControlTransformerBlock(JointTransformerBlock):
-    def __init__(
-        self,
-        layer_id: int,
-        dim: int,
-        n_heads: int,
-        n_kv_heads: int,
-        multiple_of: int,
-        ffn_dim_multiplier: float,
-        norm_eps: float,
-        qk_norm: bool,
-        modulation=True,
-        block_id=0,
-        operation_settings=None,
-    ):
-        super().__init__(layer_id, dim, n_heads, n_kv_heads, multiple_of, ffn_dim_multiplier, norm_eps, qk_norm, modulation, z_image_modulation=True, operation_settings=operation_settings)
-        self.block_id = block_id
-        if block_id == 0:
-            self.before_proj = operation_settings.get("operations").Linear(self.dim, self.dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.after_proj = operation_settings.get("operations").Linear(self.dim, self.dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-
-    def forward(self, c, x, **kwargs):
-        if self.block_id == 0:
-            c = self.before_proj(c) + x
-        c = super().forward(c, **kwargs)
-        c_skip = self.after_proj(c)
-        return c_skip, c
-
-class ZImage_Control(torch.nn.Module):
-    def __init__(
-        self,
-        dim: int = 3840,
-        n_heads: int = 30,
-        n_kv_heads: int = 30,
-        multiple_of: int = 256,
-        ffn_dim_multiplier: float = (8.0 / 3.0),
-        norm_eps: float = 1e-5,
-        qk_norm: bool = True,
-        n_control_layers=6,
-        control_in_dim=16,
-        additional_in_dim=0,
-        broken=False,
-        refiner_control=False,
-        dtype=None,
-        device=None,
-        operations=None,
-        **kwargs
-    ):
-        super().__init__()
-        operation_settings = {"operations": operations, "device": device, "dtype": dtype}
-
-        self.broken = broken
-        self.additional_in_dim = additional_in_dim
-        self.control_in_dim = control_in_dim
-        n_refiner_layers = 2
-        self.n_control_layers = n_control_layers
-        self.control_layers = nn.ModuleList(
-            [
-                ZImageControlTransformerBlock(
-                    i,
-                    dim,
-                    n_heads,
-                    n_kv_heads,
-                    multiple_of,
-                    ffn_dim_multiplier,
-                    norm_eps,
-                    qk_norm,
-                    block_id=i,
-                    operation_settings=operation_settings,
-                )
-                for i in range(self.n_control_layers)
-            ]
-        )
-
-        all_x_embedder = {}
-        patch_size = 2
-        f_patch_size = 1
-        x_embedder = operations.Linear(f_patch_size * patch_size * patch_size * (self.control_in_dim + self.additional_in_dim), dim, bias=True, device=device, dtype=dtype)
-        all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder
-
-        self.refiner_control = refiner_control
-
-        self.control_all_x_embedder = nn.ModuleDict(all_x_embedder)
-        if self.refiner_control:
-            self.control_noise_refiner = nn.ModuleList(
-                [
-                    ZImageControlTransformerBlock(
-                        layer_id,
-                        dim,
-                        n_heads,
-                        n_kv_heads,
-                        multiple_of,
-                        ffn_dim_multiplier,
-                        norm_eps,
-                        qk_norm,
-                        block_id=layer_id,
-                        operation_settings=operation_settings,
-                    )
-                    for layer_id in range(n_refiner_layers)
-                ]
-            )
-        else:
-            self.control_noise_refiner = nn.ModuleList(
-                [
-                    JointTransformerBlock(
-                        layer_id,
-                        dim,
-                        n_heads,
-                        n_kv_heads,
-                        multiple_of,
-                        ffn_dim_multiplier,
-                        norm_eps,
-                        qk_norm,
-                        modulation=True,
-                        z_image_modulation=True,
-                        operation_settings=operation_settings,
-                    )
-                    for layer_id in range(n_refiner_layers)
-                ]
-            )
-
-    def forward(self, cap_feats, control_context, x_freqs_cis, adaln_input):
-        patch_size = 2
-        f_patch_size = 1
-        pH = pW = patch_size
-        B, C, H, W = control_context.shape
-        control_context = self.control_all_x_embedder[f"{patch_size}-{f_patch_size}"](control_context.view(B, C, H // pH, pH, W // pW, pW).permute(0, 2, 4, 3, 5, 1).flatten(3).flatten(1, 2))
-
-        x_attn_mask = None
-        if not self.refiner_control:
-            for layer in self.control_noise_refiner:
-                control_context = layer(control_context, x_attn_mask, x_freqs_cis[:control_context.shape[0], :control_context.shape[1]], adaln_input)
-
-        return control_context
-
-    def forward_noise_refiner_block(self, layer_id, control_context, x, x_attn_mask, x_freqs_cis, adaln_input):
-        if self.refiner_control:
-            if self.broken:
-                if layer_id == 0:
-                    return self.control_layers[layer_id](control_context, x, x_mask=x_attn_mask, freqs_cis=x_freqs_cis[:control_context.shape[0], :control_context.shape[1]], adaln_input=adaln_input)
-                if layer_id > 0:
-                    out = None
-                    for i in range(1, len(self.control_layers)):
-                        o, control_context = self.control_layers[i](control_context, x, x_mask=x_attn_mask, freqs_cis=x_freqs_cis[:control_context.shape[0], :control_context.shape[1]], adaln_input=adaln_input)
-                        if out is None:
-                            out = o
-
-                    return (out, control_context)
-            else:
-                return self.control_noise_refiner[layer_id](control_context, x, x_mask=x_attn_mask, freqs_cis=x_freqs_cis[:control_context.shape[0], :control_context.shape[1]], adaln_input=adaln_input)
-        else:
-            return (None, control_context)
-
-    def forward_control_block(self, layer_id, control_context, x, x_attn_mask, x_freqs_cis, adaln_input):
-        return self.control_layers[layer_id](control_context, x, x_mask=x_attn_mask, freqs_cis=x_freqs_cis[:control_context.shape[0], :control_context.shape[1]], adaln_input=adaln_input)

comfy/ldm/lumina/model.py

@@ -11,7 +11,6 @@ import comfy.ldm.common_dit
 from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder
 from comfy.ldm.modules.attention import optimized_attention_masked
 from comfy.ldm.flux.layers import EmbedND
-from comfy.ldm.flux.math import apply_rope
 import comfy.patcher_extension
 
 
@@ -22,10 +21,6 @@ def modulate(x, scale):
 #                               Core NextDiT Model                              #
 #############################################################################
 
-def clamp_fp16(x):
-    if x.dtype == torch.float16:
-        return torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
-    return x
 
 class JointAttention(nn.Module):
     """Multi-head attention module."""
@@ -36,7 +31,6 @@ class JointAttention(nn.Module):
         n_heads: int,
         n_kv_heads: Optional[int],
         qk_norm: bool,
-        out_bias: bool = False,
         operation_settings={},
     ):
         """
@@ -65,7 +59,7 @@ class JointAttention(nn.Module):
         self.out = operation_settings.get("operations").Linear(
             n_heads * self.head_dim,
             dim,
-            bias=out_bias,
+            bias=False,
             device=operation_settings.get("device"),
             dtype=operation_settings.get("dtype"),
         )
@@ -76,6 +70,35 @@ class JointAttention(nn.Module):
         else:
             self.q_norm = self.k_norm = nn.Identity()
 
+    @staticmethod
+    def apply_rotary_emb(
+        x_in: torch.Tensor,
+        freqs_cis: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Apply rotary embeddings to input tensors using the given frequency
+        tensor.
+
+        This function applies rotary embeddings to the given query 'xq' and
+        key 'xk' tensors using the provided frequency tensor 'freqs_cis'. The
+        input tensors are reshaped as complex numbers, and the frequency tensor
+        is reshaped for broadcasting compatibility. The resulting tensors
+        contain rotary embeddings and are returned as real tensors.
+
+        Args:
+            x_in (torch.Tensor): Query or Key tensor to apply rotary embeddings.
+            freqs_cis (torch.Tensor): Precomputed frequency tensor for complex
+                exponentials.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor
+                and key tensor with rotary embeddings.
+        """
+
+        t_ = x_in.reshape(*x_in.shape[:-1], -1, 1, 2)
+        t_out = freqs_cis[..., 0] * t_[..., 0] + freqs_cis[..., 1] * t_[..., 1]
+        return t_out.reshape(*x_in.shape)
+
     def forward(
         self,
         x: torch.Tensor,
@@ -111,7 +134,8 @@ class JointAttention(nn.Module):
         xq = self.q_norm(xq)
         xk = self.k_norm(xk)
 
-        xq, xk = apply_rope(xq, xk, freqs_cis)
+        xq = JointAttention.apply_rotary_emb(xq, freqs_cis=freqs_cis)
+        xk = JointAttention.apply_rotary_emb(xk, freqs_cis=freqs_cis)
 
         n_rep = self.n_local_heads // self.n_local_kv_heads
         if n_rep >= 1:
@@ -173,7 +197,7 @@ class FeedForward(nn.Module):
 
     # @torch.compile
     def _forward_silu_gating(self, x1, x3):
-        return clamp_fp16(F.silu(x1) * x3)
+        return F.silu(x1) * x3
 
     def forward(self, x):
         return self.w2(self._forward_silu_gating(self.w1(x), self.w3(x)))
@@ -191,8 +215,6 @@ class JointTransformerBlock(nn.Module):
         norm_eps: float,
         qk_norm: bool,
         modulation=True,
-        z_image_modulation=False,
-        attn_out_bias=False,
         operation_settings={},
     ) -> None:
         """
@@ -213,10 +235,10 @@ class JointTransformerBlock(nn.Module):
         super().__init__()
         self.dim = dim
         self.head_dim = dim // n_heads
-        self.attention = JointAttention(dim, n_heads, n_kv_heads, qk_norm, out_bias=attn_out_bias, operation_settings=operation_settings)
+        self.attention = JointAttention(dim, n_heads, n_kv_heads, qk_norm, operation_settings=operation_settings)
         self.feed_forward = FeedForward(
             dim=dim,
-            hidden_dim=dim,
+            hidden_dim=4 * dim,
             multiple_of=multiple_of,
             ffn_dim_multiplier=ffn_dim_multiplier,
             operation_settings=operation_settings,
@@ -230,27 +252,16 @@ class JointTransformerBlock(nn.Module):
 
         self.modulation = modulation
         if modulation:
-            if z_image_modulation:
-                self.adaLN_modulation = nn.Sequential(
-                    operation_settings.get("operations").Linear(
-                        min(dim, 256),
-                        4 * dim,
-                        bias=True,
-                        device=operation_settings.get("device"),
-                        dtype=operation_settings.get("dtype"),
-                    ),
-                )
-            else:
-                self.adaLN_modulation = nn.Sequential(
-                    nn.SiLU(),
-                    operation_settings.get("operations").Linear(
-                        min(dim, 1024),
-                        4 * dim,
-                        bias=True,
-                        device=operation_settings.get("device"),
-                        dtype=operation_settings.get("dtype"),
-                    ),
-                )
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                operation_settings.get("operations").Linear(
+                    min(dim, 1024),
+                    4 * dim,
+                    bias=True,
+                    device=operation_settings.get("device"),
+                    dtype=operation_settings.get("dtype"),
+                ),
+            )
 
     def forward(
         self,
@@ -277,27 +288,27 @@ class JointTransformerBlock(nn.Module):
             scale_msa, gate_msa, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input).chunk(4, dim=1)
 
             x = x + gate_msa.unsqueeze(1).tanh() * self.attention_norm2(
-                clamp_fp16(self.attention(
+                self.attention(
                     modulate(self.attention_norm1(x), scale_msa),
                     x_mask,
                     freqs_cis,
                     transformer_options=transformer_options,
-                ))
+                )
             )
             x = x + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(
-                clamp_fp16(self.feed_forward(
+                self.feed_forward(
                     modulate(self.ffn_norm1(x), scale_mlp),
-                ))
+                )
             )
         else:
             assert adaln_input is None
             x = x + self.attention_norm2(
-                clamp_fp16(self.attention(
+                self.attention(
                     self.attention_norm1(x),
                     x_mask,
                     freqs_cis,
                     transformer_options=transformer_options,
-                ))
+                )
             )
             x = x + self.ffn_norm2(
                 self.feed_forward(
@@ -312,7 +323,7 @@ class FinalLayer(nn.Module):
     The final layer of NextDiT.
     """
 
-    def __init__(self, hidden_size, patch_size, out_channels, z_image_modulation=False, operation_settings={}):
+    def __init__(self, hidden_size, patch_size, out_channels, operation_settings={}):
         super().__init__()
         self.norm_final = operation_settings.get("operations").LayerNorm(
             hidden_size,
@@ -329,15 +340,10 @@ class FinalLayer(nn.Module):
             dtype=operation_settings.get("dtype"),
         )
 
-        if z_image_modulation:
-            min_mod = 256
-        else:
-            min_mod = 1024
-
         self.adaLN_modulation = nn.Sequential(
             nn.SiLU(),
             operation_settings.get("operations").Linear(
-                min(hidden_size, min_mod),
+                min(hidden_size, 1024),
                 hidden_size,
                 bias=True,
                 device=operation_settings.get("device"),
@@ -367,17 +373,12 @@ class NextDiT(nn.Module):
         n_heads: int = 32,
         n_kv_heads: Optional[int] = None,
         multiple_of: int = 256,
-        ffn_dim_multiplier: float = 4.0,
+        ffn_dim_multiplier: Optional[float] = None,
         norm_eps: float = 1e-5,
         qk_norm: bool = False,
         cap_feat_dim: int = 5120,
         axes_dims: List[int] = (16, 56, 56),
         axes_lens: List[int] = (1, 512, 512),
-        rope_theta=10000.0,
-        z_image_modulation=False,
-        time_scale=1.0,
-        pad_tokens_multiple=None,
-        clip_text_dim=None,
         image_model=None,
         device=None,
         dtype=None,
@@ -389,8 +390,6 @@ class NextDiT(nn.Module):
         self.in_channels = in_channels
         self.out_channels = in_channels
         self.patch_size = patch_size
-        self.time_scale = time_scale
-        self.pad_tokens_multiple = pad_tokens_multiple
 
         self.x_embedder = operation_settings.get("operations").Linear(
             in_features=patch_size * patch_size * in_channels,
@@ -412,7 +411,6 @@ class NextDiT(nn.Module):
                     norm_eps,
                     qk_norm,
                     modulation=True,
-                    z_image_modulation=z_image_modulation,
                     operation_settings=operation_settings,
                 )
                 for layer_id in range(n_refiner_layers)
@@ -436,7 +434,7 @@ class NextDiT(nn.Module):
             ]
         )
 
-        self.t_embedder = TimestepEmbedder(min(dim, 1024), output_size=256 if z_image_modulation else None, **operation_settings)
+        self.t_embedder = TimestepEmbedder(min(dim, 1024), **operation_settings)
         self.cap_embedder = nn.Sequential(
             operation_settings.get("operations").RMSNorm(cap_feat_dim, eps=norm_eps, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")),
             operation_settings.get("operations").Linear(
@@ -448,31 +446,6 @@ class NextDiT(nn.Module):
             ),
         )
 
-        self.clip_text_pooled_proj = None
-
-        if clip_text_dim is not None:
-            self.clip_text_dim = clip_text_dim
-            self.clip_text_pooled_proj = nn.Sequential(
-                operation_settings.get("operations").RMSNorm(clip_text_dim, eps=norm_eps, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")),
-                operation_settings.get("operations").Linear(
-                    clip_text_dim,
-                    clip_text_dim,
-                    bias=True,
-                    device=operation_settings.get("device"),
-                    dtype=operation_settings.get("dtype"),
-                ),
-            )
-            self.time_text_embed = nn.Sequential(
-                nn.SiLU(),
-                operation_settings.get("operations").Linear(
-                    min(dim, 1024) + clip_text_dim,
-                    min(dim, 1024),
-                    bias=True,
-                    device=operation_settings.get("device"),
-                    dtype=operation_settings.get("dtype"),
-                ),
-            )
-
         self.layers = nn.ModuleList(
             [
                 JointTransformerBlock(
@@ -484,24 +457,18 @@ class NextDiT(nn.Module):
                     ffn_dim_multiplier,
                     norm_eps,
                     qk_norm,
-                    z_image_modulation=z_image_modulation,
-                    attn_out_bias=False,
                     operation_settings=operation_settings,
                 )
                 for layer_id in range(n_layers)
             ]
         )
         self.norm_final = operation_settings.get("operations").RMSNorm(dim, eps=norm_eps, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.final_layer = FinalLayer(dim, patch_size, self.out_channels, z_image_modulation=z_image_modulation, operation_settings=operation_settings)
-
-        if self.pad_tokens_multiple is not None:
-            self.x_pad_token = nn.Parameter(torch.empty((1, dim), device=device, dtype=dtype))
-            self.cap_pad_token = nn.Parameter(torch.empty((1, dim), device=device, dtype=dtype))
+        self.final_layer = FinalLayer(dim, patch_size, self.out_channels, operation_settings=operation_settings)
 
         assert (dim // n_heads) == sum(axes_dims)
         self.axes_dims = axes_dims
         self.axes_lens = axes_lens
-        self.rope_embedder = EmbedND(dim=dim // n_heads, theta=rope_theta, axes_dim=axes_dims)
+        self.rope_embedder = EmbedND(dim=dim // n_heads, theta=10000.0, axes_dim=axes_dims)
         self.dim = dim
         self.n_heads = n_heads
 
@@ -536,63 +503,108 @@ class NextDiT(nn.Module):
         bsz = len(x)
         pH = pW = self.patch_size
         device = x[0].device
-        orig_x = x
+        dtype = x[0].dtype
 
-        if self.pad_tokens_multiple is not None:
-            pad_extra = (-cap_feats.shape[1]) % self.pad_tokens_multiple
-            cap_feats = torch.cat((cap_feats, self.cap_pad_token.to(device=cap_feats.device, dtype=cap_feats.dtype, copy=True).unsqueeze(0).repeat(cap_feats.shape[0], pad_extra, 1)), dim=1)
+        if cap_mask is not None:
+            l_effective_cap_len = cap_mask.sum(dim=1).tolist()
+        else:
+            l_effective_cap_len = [num_tokens] * bsz
 
-        cap_pos_ids = torch.zeros(bsz, cap_feats.shape[1], 3, dtype=torch.float32, device=device)
-        cap_pos_ids[:, :, 0] = torch.arange(cap_feats.shape[1], dtype=torch.float32, device=device) + 1.0
+        if cap_mask is not None and not torch.is_floating_point(cap_mask):
+            cap_mask = (cap_mask - 1).to(dtype) * torch.finfo(dtype).max
 
-        B, C, H, W = x.shape
-        x = self.x_embedder(x.view(B, C, H // pH, pH, W // pW, pW).permute(0, 2, 4, 3, 5, 1).flatten(3).flatten(1, 2))
+        img_sizes = [(img.size(1), img.size(2)) for img in x]
+        l_effective_img_len = [(H // pH) * (W // pW) for (H, W) in img_sizes]
 
-        rope_options = transformer_options.get("rope_options", None)
-        h_scale = 1.0
-        w_scale = 1.0
-        h_start = 0
-        w_start = 0
-        if rope_options is not None:
-            h_scale = rope_options.get("scale_y", 1.0)
-            w_scale = rope_options.get("scale_x", 1.0)
+        max_seq_len = max(
+            (cap_len+img_len for cap_len, img_len in zip(l_effective_cap_len, l_effective_img_len))
+        )
+        max_cap_len = max(l_effective_cap_len)
+        max_img_len = max(l_effective_img_len)
 
-            h_start = rope_options.get("shift_y", 0.0)
-            w_start = rope_options.get("shift_x", 0.0)
+        position_ids = torch.zeros(bsz, max_seq_len, 3, dtype=torch.float32, device=device)
 
-        H_tokens, W_tokens = H // pH, W // pW
-        x_pos_ids = torch.zeros((bsz, x.shape[1], 3), dtype=torch.float32, device=device)
-        x_pos_ids[:, :, 0] = cap_feats.shape[1] + 1
-        x_pos_ids[:, :, 1] = (torch.arange(H_tokens, dtype=torch.float32, device=device) * h_scale + h_start).view(-1, 1).repeat(1, W_tokens).flatten()
-        x_pos_ids[:, :, 2] = (torch.arange(W_tokens, dtype=torch.float32, device=device) * w_scale + w_start).view(1, -1).repeat(H_tokens, 1).flatten()
+        for i in range(bsz):
+            cap_len = l_effective_cap_len[i]
+            img_len = l_effective_img_len[i]
+            H, W = img_sizes[i]
+            H_tokens, W_tokens = H // pH, W // pW
+            assert H_tokens * W_tokens == img_len
 
-        if self.pad_tokens_multiple is not None:
-            pad_extra = (-x.shape[1]) % self.pad_tokens_multiple
-            x = torch.cat((x, self.x_pad_token.to(device=x.device, dtype=x.dtype, copy=True).unsqueeze(0).repeat(x.shape[0], pad_extra, 1)), dim=1)
-            x_pos_ids = torch.nn.functional.pad(x_pos_ids, (0, 0, 0, pad_extra))
+            rope_options = transformer_options.get("rope_options", None)
+            h_scale = 1.0
+            w_scale = 1.0
+            h_start = 0
+            w_start = 0
+            if rope_options is not None:
+                h_scale = rope_options.get("scale_y", 1.0)
+                w_scale = rope_options.get("scale_x", 1.0)
 
-        freqs_cis = self.rope_embedder(torch.cat((cap_pos_ids, x_pos_ids), dim=1)).movedim(1, 2)
+                h_start = rope_options.get("shift_y", 0.0)
+                w_start = rope_options.get("shift_x", 0.0)
 
-        patches = transformer_options.get("patches", {})
+            position_ids[i, :cap_len, 0] = torch.arange(cap_len, dtype=torch.float32, device=device)
+            position_ids[i, cap_len:cap_len+img_len, 0] = cap_len
+            row_ids = (torch.arange(H_tokens, dtype=torch.float32, device=device) * h_scale + h_start).view(-1, 1).repeat(1, W_tokens).flatten()
+            col_ids = (torch.arange(W_tokens, dtype=torch.float32, device=device) * w_scale + w_start).view(1, -1).repeat(H_tokens, 1).flatten()
+            position_ids[i, cap_len:cap_len+img_len, 1] = row_ids
+            position_ids[i, cap_len:cap_len+img_len, 2] = col_ids
+
+        freqs_cis = self.rope_embedder(position_ids).movedim(1, 2).to(dtype)
+
+        # build freqs_cis for cap and image individually
+        cap_freqs_cis_shape = list(freqs_cis.shape)
+        # cap_freqs_cis_shape[1] = max_cap_len
+        cap_freqs_cis_shape[1] = cap_feats.shape[1]
+        cap_freqs_cis = torch.zeros(*cap_freqs_cis_shape, device=device, dtype=freqs_cis.dtype)
+
+        img_freqs_cis_shape = list(freqs_cis.shape)
+        img_freqs_cis_shape[1] = max_img_len
+        img_freqs_cis = torch.zeros(*img_freqs_cis_shape, device=device, dtype=freqs_cis.dtype)
+
+        for i in range(bsz):
+            cap_len = l_effective_cap_len[i]
+            img_len = l_effective_img_len[i]
+            cap_freqs_cis[i, :cap_len] = freqs_cis[i, :cap_len]
+            img_freqs_cis[i, :img_len] = freqs_cis[i, cap_len:cap_len+img_len]
 
         # refine context
         for layer in self.context_refiner:
-            cap_feats = layer(cap_feats, cap_mask, freqs_cis[:, :cap_pos_ids.shape[1]], transformer_options=transformer_options)
+            cap_feats = layer(cap_feats, cap_mask, cap_freqs_cis, transformer_options=transformer_options)
 
-        padded_img_mask = None
-        x_input = x
-        for i, layer in enumerate(self.noise_refiner):
-            x = layer(x, padded_img_mask, freqs_cis[:, cap_pos_ids.shape[1]:], t, transformer_options=transformer_options)
-            if "noise_refiner" in patches:
-                for p in patches["noise_refiner"]:
-                    out = p({"img": x, "img_input": x_input, "txt": cap_feats, "pe": freqs_cis[:, cap_pos_ids.shape[1]:], "vec": t, "x": orig_x, "block_index": i, "transformer_options": transformer_options, "block_type": "noise_refiner"})
-                    if "img" in out:
-                        x = out["img"]
+        # refine image
+        flat_x = []
+        for i in range(bsz):
+            img = x[i]
+            C, H, W = img.size()
+            img = img.view(C, H // pH, pH, W // pW, pW).permute(1, 3, 2, 4, 0).flatten(2).flatten(0, 1)
+            flat_x.append(img)
+        x = flat_x
+        padded_img_embed = torch.zeros(bsz, max_img_len, x[0].shape[-1], device=device, dtype=x[0].dtype)
+        padded_img_mask = torch.zeros(bsz, max_img_len, dtype=dtype, device=device)
+        for i in range(bsz):
+            padded_img_embed[i, :l_effective_img_len[i]] = x[i]
+            padded_img_mask[i, l_effective_img_len[i]:] = -torch.finfo(dtype).max
+
+        padded_img_embed = self.x_embedder(padded_img_embed)
+        padded_img_mask = padded_img_mask.unsqueeze(1)
+        for layer in self.noise_refiner:
+            padded_img_embed = layer(padded_img_embed, padded_img_mask, img_freqs_cis, t, transformer_options=transformer_options)
+
+        if cap_mask is not None:
+            mask = torch.zeros(bsz, max_seq_len, dtype=dtype, device=device)
+            mask[:, :max_cap_len] = cap_mask[:, :max_cap_len]
+        else:
+            mask = None
+
+        padded_full_embed = torch.zeros(bsz, max_seq_len, self.dim, device=device, dtype=x[0].dtype)
+        for i in range(bsz):
+            cap_len = l_effective_cap_len[i]
+            img_len = l_effective_img_len[i]
+
+            padded_full_embed[i, :cap_len] = cap_feats[i, :cap_len]
+            padded_full_embed[i, cap_len:cap_len+img_len] = padded_img_embed[i, :img_len]
 
-        padded_full_embed = torch.cat((cap_feats, x), dim=1)
-        mask = None
-        img_sizes = [(H, W)] * bsz
-        l_effective_cap_len = [cap_feats.shape[1]] * bsz
         return padded_full_embed, mask, img_sizes, l_effective_cap_len, freqs_cis
 
     def forward(self, x, timesteps, context, num_tokens, attention_mask=None, **kwargs):
@@ -603,7 +615,7 @@ class NextDiT(nn.Module):
         ).execute(x, timesteps, context, num_tokens, attention_mask, **kwargs)
 
     # def forward(self, x, t, cap_feats, cap_mask):
-    def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, transformer_options={}, **kwargs):
+    def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, **kwargs):
         t = 1.0 - timesteps
         cap_feats = context
         cap_mask = attention_mask
@@ -615,41 +627,21 @@ class NextDiT(nn.Module):
         y: (N,) tensor of text tokens/features
         """
 
-        t = self.t_embedder(t * self.time_scale, dtype=x.dtype)  # (N, D)
+        t = self.t_embedder(t, dtype=x.dtype)  # (N, D)
         adaln_input = t
 
         cap_feats = self.cap_embedder(cap_feats)  # (N, L, D)  # todo check if able to batchify w.o. redundant compute
 
-        if self.clip_text_pooled_proj is not None:
-            pooled = kwargs.get("clip_text_pooled", None)
-            if pooled is not None:
-                pooled = self.clip_text_pooled_proj(pooled)
-            else:
-                pooled = torch.zeros((x.shape[0], self.clip_text_dim), device=x.device, dtype=x.dtype)
-
-            adaln_input = self.time_text_embed(torch.cat((t, pooled), dim=-1))
-
-        patches = transformer_options.get("patches", {})
+        transformer_options = kwargs.get("transformer_options", {})
         x_is_tensor = isinstance(x, torch.Tensor)
-        img, mask, img_size, cap_size, freqs_cis = self.patchify_and_embed(x, cap_feats, cap_mask, adaln_input, num_tokens, transformer_options=transformer_options)
-        freqs_cis = freqs_cis.to(img.device)
+        x, mask, img_size, cap_size, freqs_cis = self.patchify_and_embed(x, cap_feats, cap_mask, t, num_tokens, transformer_options=transformer_options)
+        freqs_cis = freqs_cis.to(x.device)
 
-        transformer_options["total_blocks"] = len(self.layers)
-        transformer_options["block_type"] = "double"
-        img_input = img
-        for i, layer in enumerate(self.layers):
-            transformer_options["block_index"] = i
-            img = layer(img, mask, freqs_cis, adaln_input, transformer_options=transformer_options)
-            if "double_block" in patches:
-                for p in patches["double_block"]:
-                    out = p({"img": img[:, cap_size[0]:], "img_input": img_input[:, cap_size[0]:], "txt": img[:, :cap_size[0]], "pe": freqs_cis[:, cap_size[0]:], "vec": adaln_input, "x": x, "block_index": i, "transformer_options": transformer_options})
-                    if "img" in out:
-                        img[:, cap_size[0]:] = out["img"]
-                    if "txt" in out:
-                        img[:, :cap_size[0]] = out["txt"]
+        for layer in self.layers:
+            x = layer(x, mask, freqs_cis, adaln_input, transformer_options=transformer_options)
 
-        img = self.final_layer(img, adaln_input)
-        img = self.unpatchify(img, img_size, cap_size, return_tensor=x_is_tensor)[:, :, :h, :w]
+        x = self.final_layer(x, adaln_input)
+        x = self.unpatchify(x, img_size, cap_size, return_tensor=x_is_tensor)[:,:,:h,:w]
 
-        return -img
+        return -x
 

comfy/ldm/models/autoencoder.py

@@ -9,8 +9,6 @@ from comfy.ldm.modules.distributions.distributions import DiagonalGaussianDistri
 from comfy.ldm.util import get_obj_from_str, instantiate_from_config
 from comfy.ldm.modules.ema import LitEma
 import comfy.ops
-from einops import rearrange
-import comfy.model_management
 
 class DiagonalGaussianRegularizer(torch.nn.Module):
     def __init__(self, sample: bool = False):
@@ -181,21 +179,6 @@ class AutoencodingEngineLegacy(AutoencodingEngine):
         self.post_quant_conv = conv_op(embed_dim, ddconfig["z_channels"], 1)
         self.embed_dim = embed_dim
 
-        if ddconfig.get("batch_norm_latent", False):
-            self.bn_eps = 1e-4
-            self.bn_momentum = 0.1
-            self.ps = [2, 2]
-            self.bn = torch.nn.BatchNorm2d(math.prod(self.ps) * ddconfig["z_channels"],
-                                           eps=self.bn_eps,
-                                           momentum=self.bn_momentum,
-                                           affine=False,
-                                           track_running_stats=True,
-                                           )
-            self.bn.eval()
-        else:
-            self.bn = None
-
-
     def get_autoencoder_params(self) -> list:
         params = super().get_autoencoder_params()
         return params
@@ -218,36 +201,11 @@ class AutoencodingEngineLegacy(AutoencodingEngine):
             z = torch.cat(z, 0)
 
         z, reg_log = self.regularization(z)
-
-        if self.bn is not None:
-            z = rearrange(z,
-                          "... c (i pi) (j pj)  -> ... (c pi pj) i j",
-                          pi=self.ps[0],
-                          pj=self.ps[1],
-                          )
-
-            z = torch.nn.functional.batch_norm(z,
-                                               comfy.model_management.cast_to(self.bn.running_mean, dtype=z.dtype, device=z.device),
-                                               comfy.model_management.cast_to(self.bn.running_var, dtype=z.dtype, device=z.device),
-                                               momentum=self.bn_momentum,
-                                               eps=self.bn_eps)
-
         if return_reg_log:
             return z, reg_log
         return z
 
     def decode(self, z: torch.Tensor, **decoder_kwargs) -> torch.Tensor:
-        if self.bn is not None:
-            s = torch.sqrt(comfy.model_management.cast_to(self.bn.running_var.view(1, -1, 1, 1), dtype=z.dtype, device=z.device) + self.bn_eps)
-            m = comfy.model_management.cast_to(self.bn.running_mean.view(1, -1, 1, 1), dtype=z.dtype, device=z.device)
-            z = z * s + m
-            z = rearrange(
-                z,
-                "... (c pi pj) i j -> ... c (i pi) (j pj)",
-                pi=self.ps[0],
-                pj=self.ps[1],
-            )
-
         if self.max_batch_size is None:
             dec = self.post_quant_conv(z)
             dec = self.decoder(dec, **decoder_kwargs)

comfy/ldm/modules/attention.py

@@ -517,7 +517,6 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
 
 @wrap_attn
 def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
-    exception_fallback = False
     if skip_reshape:
         b, _, _, dim_head = q.shape
         tensor_layout = "HND"
@@ -542,8 +541,6 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=
         out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout)
     except Exception as e:
         logging.error("Error running sage attention: {}, using pytorch attention instead.".format(e))
-        exception_fallback = True
-    if exception_fallback:
         if tensor_layout == "NHD":
             q, k, v = map(
                 lambda t: t.transpose(1, 2),

comfy/ldm/modules/diffusionmodules/mmdit.py

@@ -211,14 +211,12 @@ class TimestepEmbedder(nn.Module):
     Embeds scalar timesteps into vector representations.
     """
 
-    def __init__(self, hidden_size, frequency_embedding_size=256, output_size=None, dtype=None, device=None, operations=None):
+    def __init__(self, hidden_size, frequency_embedding_size=256, dtype=None, device=None, operations=None):
         super().__init__()
-        if output_size is None:
-            output_size = hidden_size
         self.mlp = nn.Sequential(
             operations.Linear(frequency_embedding_size, hidden_size, bias=True, dtype=dtype, device=device),
             nn.SiLU(),
-            operations.Linear(hidden_size, output_size, bias=True, dtype=dtype, device=device),
+            operations.Linear(hidden_size, hidden_size, bias=True, dtype=dtype, device=device),
         )
         self.frequency_embedding_size = frequency_embedding_size
 

comfy/ldm/modules/diffusionmodules/model.py

@@ -13,12 +13,6 @@ if model_management.xformers_enabled_vae():
     import xformers
     import xformers.ops
 
-def torch_cat_if_needed(xl, dim):
-    if len(xl) > 1:
-        return torch.cat(xl, dim)
-    else:
-        return xl[0]
-
 def get_timestep_embedding(timesteps, embedding_dim):
     """
     This matches the implementation in Denoising Diffusion Probabilistic Models:
@@ -49,37 +43,6 @@ def Normalize(in_channels, num_groups=32):
     return ops.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
 
 
-class CarriedConv3d(nn.Module):
-    def __init__(self, n_channels, out_channels, kernel_size, stride=1, dilation=1, padding=0, **kwargs):
-        super().__init__()
-        self.conv = ops.Conv3d(n_channels, out_channels, kernel_size, stride=stride, dilation=dilation, **kwargs)
-
-    def forward(self, x):
-        return self.conv(x)
-
-
-def conv_carry_causal_3d(xl, op, conv_carry_in=None, conv_carry_out=None):
-
-    x = xl[0]
-    xl.clear()
-
-    if isinstance(op, CarriedConv3d):
-        if conv_carry_in is None:
-            x = torch.nn.functional.pad(x, (1, 1, 1, 1, 2, 0), mode = 'replicate')
-        else:
-            carry_len = conv_carry_in[0].shape[2]
-            x = torch.nn.functional.pad(x, (1, 1, 1, 1, 2 - carry_len, 0), mode = 'replicate')
-            x = torch.cat([conv_carry_in.pop(0), x], dim=2)
-
-        if conv_carry_out is not None:
-            to_push = x[:, :, -2:, :, :].clone()
-            conv_carry_out.append(to_push)
-
-    out = op(x)
-
-    return out
-
-
 class VideoConv3d(nn.Module):
     def __init__(self, n_channels, out_channels, kernel_size, stride=1, dilation=1, padding_mode='replicate', padding=1, **kwargs):
         super().__init__()
@@ -126,24 +89,29 @@ class Upsample(nn.Module):
                                         stride=1,
                                         padding=1)
 
-    def forward(self, x, conv_carry_in=None, conv_carry_out=None):
+    def forward(self, x):
         scale_factor = self.scale_factor
         if isinstance(scale_factor, (int, float)):
             scale_factor = (scale_factor,) * (x.ndim - 2)
 
         if x.ndim == 5 and scale_factor[0] > 1.0:
-            results = []
-            if conv_carry_in is None:
-                first = x[:, :, :1, :, :]
-                results.append(interpolate_up(first.squeeze(2), scale_factor=scale_factor[1:]).unsqueeze(2))
-                x = x[:, :, 1:, :, :]
-            if x.shape[2] > 0:
-                results.append(interpolate_up(x, scale_factor))
-            x = torch_cat_if_needed(results, dim=2)
+            t = x.shape[2]
+            if t > 1:
+                a, b = x.split((1, t - 1), dim=2)
+                del x
+                b = interpolate_up(b, scale_factor)
+            else:
+                a = x
+
+            a = interpolate_up(a.squeeze(2), scale_factor=scale_factor[1:]).unsqueeze(2)
+            if t > 1:
+                x = torch.cat((a, b), dim=2)
+            else:
+                x = a
         else:
             x = interpolate_up(x, scale_factor)
         if self.with_conv:
-            x = conv_carry_causal_3d([x], self.conv, conv_carry_in, conv_carry_out)
+            x = self.conv(x)
         return x
 
 
@@ -159,20 +127,17 @@ class Downsample(nn.Module):
                                         stride=stride,
                                         padding=0)
 
-    def forward(self, x, conv_carry_in=None, conv_carry_out=None):
+    def forward(self, x):
         if self.with_conv:
-            if isinstance(self.conv, CarriedConv3d):
-                x = conv_carry_causal_3d([x], self.conv, conv_carry_in, conv_carry_out)
-            elif x.ndim == 4:
+            if x.ndim == 4:
                 pad = (0, 1, 0, 1)
                 mode = "constant"
                 x = torch.nn.functional.pad(x, pad, mode=mode, value=0)
-                x = self.conv(x)
             elif x.ndim == 5:
                 pad = (1, 1, 1, 1, 2, 0)
                 mode = "replicate"
                 x = torch.nn.functional.pad(x, pad, mode=mode)
-                x = self.conv(x)
+            x = self.conv(x)
         else:
             x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
         return x
@@ -218,23 +183,23 @@ class ResnetBlock(nn.Module):
                                                     stride=1,
                                                     padding=0)
 
-    def forward(self, x, temb=None, conv_carry_in=None, conv_carry_out=None):
+    def forward(self, x, temb=None):
         h = x
         h = self.norm1(h)
-        h = [ self.swish(h) ]
-        h = conv_carry_causal_3d(h, self.conv1, conv_carry_in=conv_carry_in, conv_carry_out=conv_carry_out)
+        h = self.swish(h)
+        h = self.conv1(h)
 
         if temb is not None:
             h = h + self.temb_proj(self.swish(temb))[:,:,None,None]
 
         h = self.norm2(h)
         h = self.swish(h)
-        h = [ self.dropout(h) ]
-        h = conv_carry_causal_3d(h, self.conv2, conv_carry_in=conv_carry_in, conv_carry_out=conv_carry_out)
+        h = self.dropout(h)
+        h = self.conv2(h)
 
         if self.in_channels != self.out_channels:
             if self.use_conv_shortcut:
-                x = conv_carry_causal_3d([x], self.conv_shortcut, conv_carry_in=conv_carry_in, conv_carry_out=conv_carry_out)
+                x = self.conv_shortcut(x)
             else:
                 x = self.nin_shortcut(x)
 
@@ -314,7 +279,6 @@ def pytorch_attention(q, k, v):
     orig_shape = q.shape
     B = orig_shape[0]
     C = orig_shape[1]
-    oom_fallback = False
     q, k, v = map(
         lambda t: t.view(B, 1, C, -1).transpose(2, 3).contiguous(),
         (q, k, v),
@@ -325,8 +289,6 @@ def pytorch_attention(q, k, v):
         out = out.transpose(2, 3).reshape(orig_shape)
     except model_management.OOM_EXCEPTION:
         logging.warning("scaled_dot_product_attention OOMed: switched to slice attention")
-        oom_fallback = True
-    if oom_fallback:
         out = slice_attention(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(orig_shape)
     return out
 
@@ -555,14 +517,9 @@ class Encoder(nn.Module):
         self.num_res_blocks = num_res_blocks
         self.resolution = resolution
         self.in_channels = in_channels
-        self.carried = False
 
         if conv3d:
-            if not attn_resolutions:
-                conv_op = CarriedConv3d
-                self.carried = True
-            else:
-                conv_op = VideoConv3d
+            conv_op = VideoConv3d
             mid_attn_conv_op = ops.Conv3d
         else:
             conv_op = ops.Conv2d
@@ -575,7 +532,6 @@ class Encoder(nn.Module):
                                        stride=1,
                                        padding=1)
 
-        self.time_compress = 1
         curr_res = resolution
         in_ch_mult = (1,)+tuple(ch_mult)
         self.in_ch_mult = in_ch_mult
@@ -602,15 +558,10 @@ class Encoder(nn.Module):
                 if time_compress is not None:
                     if (self.num_resolutions - 1 - i_level) > math.log2(time_compress):
                         stride = (1, 2, 2)
-                else:
-                    self.time_compress *= 2
                 down.downsample = Downsample(block_in, resamp_with_conv, stride=stride, conv_op=conv_op)
                 curr_res = curr_res // 2
             self.down.append(down)
 
-        if time_compress is not None:
-            self.time_compress = time_compress
-
         # middle
         self.mid = nn.Module()
         self.mid.block_1 = ResnetBlock(in_channels=block_in,
@@ -636,42 +587,15 @@ class Encoder(nn.Module):
     def forward(self, x):
         # timestep embedding
         temb = None
-
-        if self.carried:
-            xl = [x[:, :, :1, :, :]]
-            if x.shape[2] > self.time_compress:
-                tc = self.time_compress
-                xl += torch.split(x[:, :, 1: 1 + ((x.shape[2] - 1) // tc) * tc, :, :], tc * 2, dim = 2)
-            x = xl
-        else:
-            x = [x]
-        out = []
-
-        conv_carry_in = None
-
-        for i, x1 in enumerate(x):
-            conv_carry_out = []
-            if i == len(x) - 1:
-                conv_carry_out = None
-
-            # downsampling
-            x1 = [ x1 ]
-            h1 = conv_carry_causal_3d(x1, self.conv_in, conv_carry_in, conv_carry_out)
-
-            for i_level in range(self.num_resolutions):
-                for i_block in range(self.num_res_blocks):
-                    h1 = self.down[i_level].block[i_block](h1, temb, conv_carry_in, conv_carry_out)
-                    if len(self.down[i_level].attn) > 0:
-                        assert i == 0 #carried should not happen if attn exists
-                        h1 = self.down[i_level].attn[i_block](h1)
-                if i_level != self.num_resolutions-1:
-                    h1 = self.down[i_level].downsample(h1, conv_carry_in, conv_carry_out)
-
-            out.append(h1)
-            conv_carry_in = conv_carry_out
-
-        h = torch_cat_if_needed(out, dim=2)
-        del out
+        # downsampling
+        h = self.conv_in(x)
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](h, temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+            if i_level != self.num_resolutions-1:
+                h = self.down[i_level].downsample(h)
 
         # middle
         h = self.mid.block_1(h, temb)
@@ -680,15 +604,15 @@ class Encoder(nn.Module):
 
         # end
         h = self.norm_out(h)
-        h = [ nonlinearity(h) ]
-        h = conv_carry_causal_3d(h, self.conv_out)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
         return h
 
 
 class Decoder(nn.Module):
     def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
                  attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, tanh_out=False, use_linear_attn=False,
+                 resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
                  conv_out_op=ops.Conv2d,
                  resnet_op=ResnetBlock,
                  attn_op=AttnBlock,
@@ -702,18 +626,12 @@ class Decoder(nn.Module):
         self.num_res_blocks = num_res_blocks
         self.resolution = resolution
         self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
         self.tanh_out = tanh_out
-        self.carried = False
 
         if conv3d:
-            if not attn_resolutions and resnet_op == ResnetBlock:
-                conv_op = CarriedConv3d
-                conv_out_op = CarriedConv3d
-                self.carried = True
-            else:
-                conv_op = VideoConv3d
-                conv_out_op = VideoConv3d
-
+            conv_op = VideoConv3d
+            conv_out_op = VideoConv3d
             mid_attn_conv_op = ops.Conv3d
         else:
             conv_op = ops.Conv2d
@@ -788,43 +706,29 @@ class Decoder(nn.Module):
         temb = None
 
         # z to block_in
-        h = conv_carry_causal_3d([z], self.conv_in)
+        h = self.conv_in(z)
 
         # middle
         h = self.mid.block_1(h, temb, **kwargs)
         h = self.mid.attn_1(h, **kwargs)
         h = self.mid.block_2(h, temb, **kwargs)
 
-        if self.carried:
-            h = torch.split(h, 2, dim=2)
-        else:
-            h = [ h ]
-        out = []
-
-        conv_carry_in = None
-
         # upsampling
-        for i, h1 in enumerate(h):
-            conv_carry_out = []
-            if i == len(h) - 1:
-                conv_carry_out = None
-            for i_level in reversed(range(self.num_resolutions)):
-                for i_block in range(self.num_res_blocks+1):
-                    h1 = self.up[i_level].block[i_block](h1, temb, conv_carry_in, conv_carry_out, **kwargs)
-                    if len(self.up[i_level].attn) > 0:
-                        assert i == 0 #carried should not happen if attn exists
-                        h1 = self.up[i_level].attn[i_block](h1, **kwargs)
-                if i_level != 0:
-                    h1 = self.up[i_level].upsample(h1, conv_carry_in, conv_carry_out)
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](h, temb, **kwargs)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h, **kwargs)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
 
-            h1 = self.norm_out(h1)
-            h1 = [ nonlinearity(h1) ]
-            h1 = conv_carry_causal_3d(h1, self.conv_out, conv_carry_in, conv_carry_out)
-            if self.tanh_out:
-                h1 = torch.tanh(h1)
-            out.append(h1)
-            conv_carry_in = conv_carry_out
+        # end
+        if self.give_pre_end:
+            return h
 
-        out = torch_cat_if_needed(out, dim=2)
-
-        return out
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h, **kwargs)
+        if self.tanh_out:
+            h = torch.tanh(h)
+        return h

comfy/ldm/qwen_image/model.py

@@ -61,7 +61,7 @@ def apply_rotary_emb(x, freqs_cis):
 
 
 class QwenTimestepProjEmbeddings(nn.Module):
-    def __init__(self, embedding_dim, pooled_projection_dim, use_additional_t_cond=False, dtype=None, device=None, operations=None):
+    def __init__(self, embedding_dim, pooled_projection_dim, dtype=None, device=None, operations=None):
         super().__init__()
         self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000)
         self.timestep_embedder = TimestepEmbedding(
@@ -72,19 +72,9 @@ class QwenTimestepProjEmbeddings(nn.Module):
             operations=operations
         )
 
-        self.use_additional_t_cond = use_additional_t_cond
-        if self.use_additional_t_cond:
-            self.addition_t_embedding = operations.Embedding(2, embedding_dim, device=device, dtype=dtype)
-
-    def forward(self, timestep, hidden_states, addition_t_cond=None):
+    def forward(self, timestep, hidden_states):
         timesteps_proj = self.time_proj(timestep)
         timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype))
-
-        if self.use_additional_t_cond:
-            if addition_t_cond is None:
-                addition_t_cond = torch.zeros((timesteps_emb.shape[0]), device=timesteps_emb.device, dtype=torch.long)
-            timesteps_emb += self.addition_t_embedding(addition_t_cond, out_dtype=timesteps_emb.dtype)
-
         return timesteps_emb
 
 
@@ -228,24 +218,9 @@ class QwenImageTransformerBlock(nn.Module):
             operations=operations,
         )
 
-    def _apply_gate(self, x, y, gate, timestep_zero_index=None):
-        if timestep_zero_index is not None:
-            return y + torch.cat((x[:, :timestep_zero_index] * gate[0], x[:, timestep_zero_index:] * gate[1]), dim=1)
-        else:
-            return torch.addcmul(y, gate, x)
-
-    def _modulate(self, x: torch.Tensor, mod_params: torch.Tensor, timestep_zero_index=None) -> Tuple[torch.Tensor, torch.Tensor]:
+    def _modulate(self, x: torch.Tensor, mod_params: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         shift, scale, gate = torch.chunk(mod_params, 3, dim=-1)
-        if timestep_zero_index is not None:
-            actual_batch = shift.size(0) // 2
-            shift, shift_0 = shift[:actual_batch], shift[actual_batch:]
-            scale, scale_0 = scale[:actual_batch], scale[actual_batch:]
-            gate, gate_0 = gate[:actual_batch], gate[actual_batch:]
-            reg = torch.addcmul(shift.unsqueeze(1), x[:, :timestep_zero_index], 1 + scale.unsqueeze(1))
-            zero = torch.addcmul(shift_0.unsqueeze(1), x[:, timestep_zero_index:], 1 + scale_0.unsqueeze(1))
-            return torch.cat((reg, zero), dim=1), (gate.unsqueeze(1), gate_0.unsqueeze(1))
-        else:
-            return torch.addcmul(shift.unsqueeze(1), x, 1 + scale.unsqueeze(1)), gate.unsqueeze(1)
+        return torch.addcmul(shift.unsqueeze(1), x, 1 + scale.unsqueeze(1)), gate.unsqueeze(1)
 
     def forward(
         self,
@@ -254,22 +229,17 @@ class QwenImageTransformerBlock(nn.Module):
         encoder_hidden_states_mask: torch.Tensor,
         temb: torch.Tensor,
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        timestep_zero_index=None,
         transformer_options={},
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         img_mod_params = self.img_mod(temb)
-
-        if timestep_zero_index is not None:
-            temb = temb.chunk(2, dim=0)[0]
-
         txt_mod_params = self.txt_mod(temb)
         img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)
         txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)
 
-        img_modulated, img_gate1 = self._modulate(self.img_norm1(hidden_states), img_mod1, timestep_zero_index)
-        del img_mod1
-        txt_modulated, txt_gate1 = self._modulate(self.txt_norm1(encoder_hidden_states), txt_mod1)
-        del txt_mod1
+        img_normed = self.img_norm1(hidden_states)
+        img_modulated, img_gate1 = self._modulate(img_normed, img_mod1)
+        txt_normed = self.txt_norm1(encoder_hidden_states)
+        txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1)
 
         img_attn_output, txt_attn_output = self.attn(
             hidden_states=img_modulated,
@@ -278,20 +248,16 @@ class QwenImageTransformerBlock(nn.Module):
             image_rotary_emb=image_rotary_emb,
             transformer_options=transformer_options,
         )
-        del img_modulated
-        del txt_modulated
 
-        hidden_states = self._apply_gate(img_attn_output, hidden_states, img_gate1, timestep_zero_index)
+        hidden_states = hidden_states + img_gate1 * img_attn_output
         encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
-        del img_attn_output
-        del txt_attn_output
-        del img_gate1
-        del txt_gate1
 
-        img_modulated2, img_gate2 = self._modulate(self.img_norm2(hidden_states), img_mod2, timestep_zero_index)
-        hidden_states = self._apply_gate(self.img_mlp(img_modulated2), hidden_states, img_gate2, timestep_zero_index)
+        img_normed2 = self.img_norm2(hidden_states)
+        img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
+        hidden_states = torch.addcmul(hidden_states, img_gate2, self.img_mlp(img_modulated2))
 
-        txt_modulated2, txt_gate2 = self._modulate(self.txt_norm2(encoder_hidden_states), txt_mod2)
+        txt_normed2 = self.txt_norm2(encoder_hidden_states)
+        txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
         encoder_hidden_states = torch.addcmul(encoder_hidden_states, txt_gate2, self.txt_mlp(txt_modulated2))
 
         return encoder_hidden_states, hidden_states
@@ -330,11 +296,10 @@ class QwenImageTransformer2DModel(nn.Module):
         num_attention_heads: int = 24,
         joint_attention_dim: int = 3584,
         pooled_projection_dim: int = 768,
+        guidance_embeds: bool = False,
         axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
-        default_ref_method="index",
         image_model=None,
         final_layer=True,
-        use_additional_t_cond=False,
         dtype=None,
         device=None,
         operations=None,
@@ -345,14 +310,12 @@ class QwenImageTransformer2DModel(nn.Module):
         self.in_channels = in_channels
         self.out_channels = out_channels or in_channels
         self.inner_dim = num_attention_heads * attention_head_dim
-        self.default_ref_method = default_ref_method
 
         self.pe_embedder = EmbedND(dim=attention_head_dim, theta=10000, axes_dim=list(axes_dims_rope))
 
         self.time_text_embed = QwenTimestepProjEmbeddings(
             embedding_dim=self.inner_dim,
             pooled_projection_dim=pooled_projection_dim,
-            use_additional_t_cond=use_additional_t_cond,
             dtype=dtype,
             device=device,
             operations=operations
@@ -374,9 +337,6 @@ class QwenImageTransformer2DModel(nn.Module):
             for _ in range(num_layers)
         ])
 
-        if self.default_ref_method == "index_timestep_zero":
-            self.register_buffer("__index_timestep_zero__", torch.tensor([]))
-
         if final_layer:
             self.norm_out = LastLayer(self.inner_dim, self.inner_dim, dtype=dtype, device=device, operations=operations)
             self.proj_out = operations.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True, dtype=dtype, device=device)
@@ -386,33 +346,27 @@ class QwenImageTransformer2DModel(nn.Module):
         patch_size = self.patch_size
         hidden_states = comfy.ldm.common_dit.pad_to_patch_size(x, (1, self.patch_size, self.patch_size))
         orig_shape = hidden_states.shape
-        hidden_states = hidden_states.view(orig_shape[0], orig_shape[1], orig_shape[-3], orig_shape[-2] // 2, 2, orig_shape[-1] // 2, 2)
-        hidden_states = hidden_states.permute(0, 2, 3, 5, 1, 4, 6)
-        hidden_states = hidden_states.reshape(orig_shape[0], orig_shape[-3] * (orig_shape[-2] // 2) * (orig_shape[-1] // 2), orig_shape[1] * 4)
-        t_len = t
+        hidden_states = hidden_states.view(orig_shape[0], orig_shape[1], orig_shape[-2] // 2, 2, orig_shape[-1] // 2, 2)
+        hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5)
+        hidden_states = hidden_states.reshape(orig_shape[0], (orig_shape[-2] // 2) * (orig_shape[-1] // 2), orig_shape[1] * 4)
         h_len = ((h + (patch_size // 2)) // patch_size)
         w_len = ((w + (patch_size // 2)) // patch_size)
 
         h_offset = ((h_offset + (patch_size // 2)) // patch_size)
         w_offset = ((w_offset + (patch_size // 2)) // patch_size)
 
-        img_ids = torch.zeros((t_len, h_len, w_len, 3), device=x.device)
+        img_ids = torch.zeros((h_len, w_len, 3), device=x.device)
+        img_ids[:, :, 0] = img_ids[:, :, 1] + index
+        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(h_offset, h_len - 1 + h_offset, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1) - (h_len // 2)
+        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(w_offset, w_len - 1 + w_offset, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0) - (w_len // 2)
+        return hidden_states, repeat(img_ids, "h w c -> b (h w) c", b=bs), orig_shape
 
-        if t_len > 1:
-            img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.linspace(0, t_len - 1, steps=t_len, device=x.device, dtype=x.dtype).unsqueeze(1).unsqueeze(1)
-        else:
-            img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + index
-
-        img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.linspace(h_offset, h_len - 1 + h_offset, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1).unsqueeze(0) - (h_len // 2)
-        img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.linspace(w_offset, w_len - 1 + w_offset, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0).unsqueeze(0) - (w_len // 2)
-        return hidden_states, repeat(img_ids, "t h w c -> b (t h w) c", b=bs), orig_shape
-
-    def forward(self, x, timestep, context, attention_mask=None, ref_latents=None, additional_t_cond=None, transformer_options={}, **kwargs):
+    def forward(self, x, timestep, context, attention_mask=None, guidance=None, ref_latents=None, transformer_options={}, **kwargs):
         return comfy.patcher_extension.WrapperExecutor.new_class_executor(
             self._forward,
             self,
             comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
-        ).execute(x, timestep, context, attention_mask, ref_latents, additional_t_cond, transformer_options, **kwargs)
+        ).execute(x, timestep, context, attention_mask, guidance, ref_latents, transformer_options, **kwargs)
 
     def _forward(
         self,
@@ -420,8 +374,8 @@ class QwenImageTransformer2DModel(nn.Module):
         timesteps,
         context,
         attention_mask=None,
+        guidance: torch.Tensor = None,
         ref_latents=None,
-        additional_t_cond=None,
         transformer_options={},
         control=None,
         **kwargs
@@ -433,24 +387,16 @@ class QwenImageTransformer2DModel(nn.Module):
         hidden_states, img_ids, orig_shape = self.process_img(x)
         num_embeds = hidden_states.shape[1]
 
-        timestep_zero_index = None
         if ref_latents is not None:
             h = 0
             w = 0
             index = 0
-            ref_method = kwargs.get("ref_latents_method", self.default_ref_method)
-            index_ref_method = (ref_method == "index") or (ref_method == "index_timestep_zero")
-            negative_ref_method = ref_method == "negative_index"
-            timestep_zero = ref_method == "index_timestep_zero"
+            index_ref_method = kwargs.get("ref_latents_method", "index") == "index"
             for ref in ref_latents:
                 if index_ref_method:
                     index += 1
                     h_offset = 0
                     w_offset = 0
-                elif negative_ref_method:
-                    index -= 1
-                    h_offset = 0
-                    w_offset = 0
                 else:
                     index = 1
                     h_offset = 0
@@ -465,10 +411,6 @@ class QwenImageTransformer2DModel(nn.Module):
                 kontext, kontext_ids, _ = self.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset)
                 hidden_states = torch.cat([hidden_states, kontext], dim=1)
                 img_ids = torch.cat([img_ids, kontext_ids], dim=1)
-            if timestep_zero:
-                if index > 0:
-                    timestep = torch.cat([timestep, timestep * 0], dim=0)
-                    timestep_zero_index = num_embeds
 
         txt_start = round(max(((x.shape[-1] + (self.patch_size // 2)) // self.patch_size) // 2, ((x.shape[-2] + (self.patch_size // 2)) // self.patch_size) // 2))
         txt_ids = torch.arange(txt_start, txt_start + context.shape[1], device=x.device).reshape(1, -1, 1).repeat(x.shape[0], 1, 3)
@@ -480,20 +422,24 @@ class QwenImageTransformer2DModel(nn.Module):
         encoder_hidden_states = self.txt_norm(encoder_hidden_states)
         encoder_hidden_states = self.txt_in(encoder_hidden_states)
 
-        temb = self.time_text_embed(timestep, hidden_states, additional_t_cond)
+        if guidance is not None:
+            guidance = guidance * 1000
+
+        temb = (
+            self.time_text_embed(timestep, hidden_states)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, hidden_states)
+        )
 
         patches_replace = transformer_options.get("patches_replace", {})
         patches = transformer_options.get("patches", {})
         blocks_replace = patches_replace.get("dit", {})
 
-        transformer_options["total_blocks"] = len(self.transformer_blocks)
-        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.transformer_blocks):
-            transformer_options["block_index"] = i
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
-                    out["txt"], out["img"] = block(hidden_states=args["img"], encoder_hidden_states=args["txt"], encoder_hidden_states_mask=encoder_hidden_states_mask, temb=args["vec"], image_rotary_emb=args["pe"], timestep_zero_index=timestep_zero_index, transformer_options=args["transformer_options"])
+                    out["txt"], out["img"] = block(hidden_states=args["img"], encoder_hidden_states=args["txt"], encoder_hidden_states_mask=encoder_hidden_states_mask, temb=args["vec"], image_rotary_emb=args["pe"], transformer_options=args["transformer_options"])
                     return out
                 out = blocks_replace[("double_block", i)]({"img": hidden_states, "txt": encoder_hidden_states, "vec": temb, "pe": image_rotary_emb, "transformer_options": transformer_options}, {"original_block": block_wrap})
                 hidden_states = out["img"]
@@ -505,7 +451,6 @@ class QwenImageTransformer2DModel(nn.Module):
                     encoder_hidden_states_mask=encoder_hidden_states_mask,
                     temb=temb,
                     image_rotary_emb=image_rotary_emb,
-                    timestep_zero_index=timestep_zero_index,
                     transformer_options=transformer_options,
                 )
 
@@ -522,12 +467,9 @@ class QwenImageTransformer2DModel(nn.Module):
                     if add is not None:
                         hidden_states[:, :add.shape[1]] += add
 
-        if timestep_zero_index is not None:
-            temb = temb.chunk(2, dim=0)[0]
-
         hidden_states = self.norm_out(hidden_states, temb)
         hidden_states = self.proj_out(hidden_states)
 
-        hidden_states = hidden_states[:, :num_embeds].view(orig_shape[0], orig_shape[-3], orig_shape[-2] // 2, orig_shape[-1] // 2, orig_shape[1], 2, 2)
-        hidden_states = hidden_states.permute(0, 4, 1, 2, 5, 3, 6)
+        hidden_states = hidden_states[:, :num_embeds].view(orig_shape[0], orig_shape[-2] // 2, orig_shape[-1] // 2, orig_shape[1], 2, 2)
+        hidden_states = hidden_states.permute(0, 3, 1, 4, 2, 5)
         return hidden_states.reshape(orig_shape)[:, :, :, :x.shape[-2], :x.shape[-1]]

comfy/ldm/wan/model.py

@@ -568,10 +568,7 @@ class WanModel(torch.nn.Module):
 
         patches_replace = transformer_options.get("patches_replace", {})
         blocks_replace = patches_replace.get("dit", {})
-        transformer_options["total_blocks"] = len(self.blocks)
-        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.blocks):
-            transformer_options["block_index"] = i
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -766,10 +763,7 @@ class VaceWanModel(WanModel):
 
         patches_replace = transformer_options.get("patches_replace", {})
         blocks_replace = patches_replace.get("dit", {})
-        transformer_options["total_blocks"] = len(self.blocks)
-        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.blocks):
-            transformer_options["block_index"] = i
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -868,10 +862,7 @@ class CameraWanModel(WanModel):
 
         patches_replace = transformer_options.get("patches_replace", {})
         blocks_replace = patches_replace.get("dit", {})
-        transformer_options["total_blocks"] = len(self.blocks)
-        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.blocks):
-            transformer_options["block_index"] = i
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -1335,19 +1326,16 @@ class WanModel_S2V(WanModel):
 
         patches_replace = transformer_options.get("patches_replace", {})
         blocks_replace = patches_replace.get("dit", {})
-        transformer_options["total_blocks"] = len(self.blocks)
-        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.blocks):
-            transformer_options["block_index"] = i
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
-                    out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], transformer_options=args["transformer_options"])
+                    out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"])
                     return out
-                out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
+                out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs}, {"original_block": block_wrap})
                 x = out["img"]
             else:
-                x = block(x, e=e0, freqs=freqs, context=context, transformer_options=transformer_options)
+                x = block(x, e=e0, freqs=freqs, context=context)
             if audio_emb is not None:
                 x = self.audio_injector(x, i, audio_emb, audio_emb_global, seq_len)
         # head
@@ -1586,10 +1574,7 @@ class HumoWanModel(WanModel):
 
         patches_replace = transformer_options.get("patches_replace", {})
         blocks_replace = patches_replace.get("dit", {})
-        transformer_options["total_blocks"] = len(self.blocks)
-        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.blocks):
-            transformer_options["block_index"] = i
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}

comfy/ldm/wan/model_animate.py

@@ -523,10 +523,7 @@ class AnimateWanModel(WanModel):
 
         patches_replace = transformer_options.get("patches_replace", {})
         blocks_replace = patches_replace.get("dit", {})
-        transformer_options["total_blocks"] = len(self.blocks)
-        transformer_options["block_type"] = "double"
         for i, block in enumerate(self.blocks):
-            transformer_options["block_index"] = i
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}

comfy/ldm/wan/vae.py

@@ -227,7 +227,6 @@ class Encoder3d(nn.Module):
     def __init__(self,
                  dim=128,
                  z_dim=4,
-                 input_channels=3,
                  dim_mult=[1, 2, 4, 4],
                  num_res_blocks=2,
                  attn_scales=[],
@@ -246,7 +245,7 @@ class Encoder3d(nn.Module):
         scale = 1.0
 
         # init block
-        self.conv1 = CausalConv3d(input_channels, dims[0], 3, padding=1)
+        self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)
 
         # downsample blocks
         downsamples = []
@@ -332,7 +331,6 @@ class Decoder3d(nn.Module):
     def __init__(self,
                  dim=128,
                  z_dim=4,
-                 output_channels=3,
                  dim_mult=[1, 2, 4, 4],
                  num_res_blocks=2,
                  attn_scales=[],
@@ -380,7 +378,7 @@ class Decoder3d(nn.Module):
         # output blocks
         self.head = nn.Sequential(
             RMS_norm(out_dim, images=False), nn.SiLU(),
-            CausalConv3d(out_dim, output_channels, 3, padding=1))
+            CausalConv3d(out_dim, 3, 3, padding=1))
 
     def forward(self, x, feat_cache=None, feat_idx=[0]):
         ## conv1
@@ -451,7 +449,6 @@ class WanVAE(nn.Module):
                  num_res_blocks=2,
                  attn_scales=[],
                  temperal_downsample=[True, True, False],
-                 image_channels=3,
                  dropout=0.0):
         super().__init__()
         self.dim = dim
@@ -463,11 +460,11 @@ class WanVAE(nn.Module):
         self.temperal_upsample = temperal_downsample[::-1]
 
         # modules
-        self.encoder = Encoder3d(dim, z_dim * 2, image_channels, dim_mult, num_res_blocks,
+        self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,
                                  attn_scales, self.temperal_downsample, dropout)
         self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
         self.conv2 = CausalConv3d(z_dim, z_dim, 1)
-        self.decoder = Decoder3d(dim, z_dim, image_channels, dim_mult, num_res_blocks,
+        self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks,
                                  attn_scales, self.temperal_upsample, dropout)
 
     def encode(self, x):

comfy/lora.py

@@ -313,23 +313,6 @@ def model_lora_keys_unet(model, key_map={}):
                 key_map["transformer.{}".format(key_lora)] = k
                 key_map["lycoris_{}".format(key_lora.replace(".", "_"))] = k #SimpleTuner lycoris format
 
-    if isinstance(model, comfy.model_base.Lumina2):
-        diffusers_keys = comfy.utils.z_image_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 = k[:-len(".weight")]
-                key_map["diffusion_model.{}".format(key_lora)] = to
-                key_map["transformer.{}".format(key_lora)] = to
-                key_map["lycoris_{}".format(key_lora.replace(".", "_"))] = to
-
-    if isinstance(model, comfy.model_base.Kandinsky5):
-        for k in sdk:
-            if k.startswith("diffusion_model.") and k.endswith(".weight"):
-                key_lora = k[len("diffusion_model."):-len(".weight")]
-                key_map["{}".format(key_lora)] = k
-                key_map["transformer.{}".format(key_lora)] = k
-
     return key_map
 
 

comfy/model_base.py

@@ -47,7 +47,6 @@ import comfy.ldm.chroma_radiance.model
 import comfy.ldm.ace.model
 import comfy.ldm.omnigen.omnigen2
 import comfy.ldm.qwen_image.model
-import comfy.ldm.kandinsky5.model
 
 import comfy.model_management
 import comfy.patcher_extension
@@ -135,7 +134,7 @@ class BaseModel(torch.nn.Module):
         if not unet_config.get("disable_unet_model_creation", False):
             if model_config.custom_operations is None:
                 fp8 = model_config.optimizations.get("fp8", False)
-                operations = comfy.ops.pick_operations(unet_config.get("dtype", None), self.manual_cast_dtype, fp8_optimizations=fp8, model_config=model_config)
+                operations = comfy.ops.pick_operations(unet_config.get("dtype", None), self.manual_cast_dtype, fp8_optimizations=fp8, scaled_fp8=model_config.scaled_fp8, model_config=model_config)
             else:
                 operations = model_config.custom_operations
             self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
@@ -330,6 +329,18 @@ class BaseModel(torch.nn.Module):
             extra_sds.append(self.model_config.process_clip_vision_state_dict_for_saving(clip_vision_state_dict))
 
         unet_state_dict = self.diffusion_model.state_dict()
+
+        if self.model_config.scaled_fp8 is not None:
+            unet_state_dict["scaled_fp8"] = torch.tensor([], dtype=self.model_config.scaled_fp8)
+
+        # Save mixed precision metadata
+        if hasattr(self.model_config, 'layer_quant_config') and self.model_config.layer_quant_config:
+            metadata = {
+                "format_version": "1.0",
+                "layers": self.model_config.layer_quant_config
+            }
+            unet_state_dict["_quantization_metadata"] = metadata
+
         unet_state_dict = self.model_config.process_unet_state_dict_for_saving(unet_state_dict)
 
         if self.model_type == ModelType.V_PREDICTION:
@@ -887,13 +898,12 @@ class Flux(BaseModel):
         attention_mask = kwargs.get("attention_mask", None)
         if attention_mask is not None:
             shape = kwargs["noise"].shape
-            mask_ref_size = kwargs.get("attention_mask_img_shape", None)
-            if mask_ref_size is not None:
-                # the model will pad to the patch size, and then divide
-                # essentially dividing and rounding up
-                (h_tok, w_tok) = (math.ceil(shape[2] / self.diffusion_model.patch_size), math.ceil(shape[3] / self.diffusion_model.patch_size))
-                attention_mask = utils.upscale_dit_mask(attention_mask, mask_ref_size, (h_tok, w_tok))
-                out['attention_mask'] = comfy.conds.CONDRegular(attention_mask)
+            mask_ref_size = kwargs["attention_mask_img_shape"]
+            # the model will pad to the patch size, and then divide
+            # essentially dividing and rounding up
+            (h_tok, w_tok) = (math.ceil(shape[2] / self.diffusion_model.patch_size), math.ceil(shape[3] / self.diffusion_model.patch_size))
+            attention_mask = utils.upscale_dit_mask(attention_mask, mask_ref_size, (h_tok, w_tok))
+            out['attention_mask'] = comfy.conds.CONDRegular(attention_mask)
 
         guidance = kwargs.get("guidance", 3.5)
         if guidance is not None:
@@ -915,19 +925,9 @@ class Flux(BaseModel):
         out = {}
         ref_latents = kwargs.get("reference_latents", None)
         if ref_latents is not None:
-            out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()[2:]), ref_latents))])
+            out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
         return out
 
-class Flux2(Flux):
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-        cross_attn = kwargs.get("cross_attn", None)
-        if cross_attn is not None:
-            target_text_len = 512
-            if cross_attn.shape[1] < target_text_len:
-                cross_attn = torch.nn.functional.pad(cross_attn, (0, 0, target_text_len - cross_attn.shape[1], 0))
-            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
-        return out
 
 class GenmoMochi(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
@@ -1103,17 +1103,9 @@ class Lumina2(BaseModel):
             if torch.numel(attention_mask) != attention_mask.sum():
                 out['attention_mask'] = comfy.conds.CONDRegular(attention_mask)
             out['num_tokens'] = comfy.conds.CONDConstant(max(1, torch.sum(attention_mask).item()))
-
         cross_attn = kwargs.get("cross_attn", None)
         if cross_attn is not None:
             out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
-            if 'num_tokens' not in out:
-                out['num_tokens'] = comfy.conds.CONDConstant(cross_attn.shape[1])
-
-        clip_text_pooled = kwargs.get("pooled_output", None)  # NewBie
-        if clip_text_pooled is not None:
-            out['clip_text_pooled'] = comfy.conds.CONDRegular(clip_text_pooled)
-
         return out
 
 class WAN21(BaseModel):
@@ -1544,140 +1536,3 @@ class HunyuanImage21Refiner(HunyuanImage21):
         out = super().extra_conds(**kwargs)
         out['disable_time_r'] = comfy.conds.CONDConstant(True)
         return out
-
-class HunyuanVideo15(HunyuanVideo):
-    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
-        super().__init__(model_config, model_type, device=device)
-
-    def concat_cond(self, **kwargs):
-        noise = kwargs.get("noise", None)
-        extra_channels = self.diffusion_model.img_in.proj.weight.shape[1] - noise.shape[1] - 1 #noise 32 img cond 32 + mask 1
-        if extra_channels == 0:
-            return None
-
-        image = kwargs.get("concat_latent_image", None)
-        device = kwargs["device"]
-
-        if image is None:
-            shape_image = list(noise.shape)
-            shape_image[1] = extra_channels
-            image = torch.zeros(shape_image, dtype=noise.dtype, layout=noise.layout, device=noise.device)
-        else:
-            latent_dim = self.latent_format.latent_channels
-            image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
-            for i in range(0, image.shape[1], latent_dim):
-                image[:, i: i + latent_dim] = self.process_latent_in(image[:, i: i + latent_dim])
-            image = utils.resize_to_batch_size(image, noise.shape[0])
-
-        mask = kwargs.get("concat_mask", kwargs.get("denoise_mask", None))
-        if mask is None:
-            mask = torch.zeros_like(noise)[:, :1]
-        else:
-            mask = 1.0 - mask
-            mask = utils.common_upscale(mask.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
-            if mask.shape[-3] < noise.shape[-3]:
-                mask = torch.nn.functional.pad(mask, (0, 0, 0, 0, 0, noise.shape[-3] - mask.shape[-3]), mode='constant', value=0)
-            mask = utils.resize_to_batch_size(mask, noise.shape[0])
-
-        return torch.cat((image, mask), dim=1)
-
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-        attention_mask = kwargs.get("attention_mask", None)
-        if attention_mask is not None:
-            if torch.numel(attention_mask) != attention_mask.sum():
-                out['attention_mask'] = comfy.conds.CONDRegular(attention_mask)
-        cross_attn = kwargs.get("cross_attn", None)
-        if cross_attn is not None:
-            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
-
-        conditioning_byt5small = kwargs.get("conditioning_byt5small", None)
-        if conditioning_byt5small is not None:
-            out['txt_byt5'] = comfy.conds.CONDRegular(conditioning_byt5small)
-
-        guidance = kwargs.get("guidance", 6.0)
-        if guidance is not None:
-            out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
-
-        clip_vision_output = kwargs.get("clip_vision_output", None)
-        if clip_vision_output is not None:
-            out['clip_fea'] = comfy.conds.CONDRegular(clip_vision_output.last_hidden_state)
-
-        return out
-
-class HunyuanVideo15_SR_Distilled(HunyuanVideo15):
-    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
-        super().__init__(model_config, model_type, device=device)
-
-    def concat_cond(self, **kwargs):
-        noise = kwargs.get("noise", None)
-        image = kwargs.get("concat_latent_image", None)
-        noise_augmentation = kwargs.get("noise_augmentation", 0.0)
-        device = kwargs["device"]
-
-        if image is None:
-            image = torch.zeros([noise.shape[0], noise.shape[1] * 2 + 2, noise.shape[-3], noise.shape[-2], noise.shape[-1]], device=comfy.model_management.intermediate_device())
-        else:
-            image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
-            #image = self.process_latent_in(image) # scaling wasn't applied in reference code
-            image = utils.resize_to_batch_size(image, noise.shape[0])
-            lq_image_slice = slice(noise.shape[1] + 1, 2 * noise.shape[1] + 1)
-            if noise_augmentation > 0:
-                generator = torch.Generator(device="cpu")
-                generator.manual_seed(kwargs.get("seed", 0) - 10)
-                noise = torch.randn(image[:, lq_image_slice].shape, generator=generator, dtype=image.dtype, device="cpu").to(image.device)
-                image[:, lq_image_slice] = noise_augmentation * noise + min(1.0 - noise_augmentation, 0.75) * image[:, lq_image_slice]
-            else:
-                image[:, lq_image_slice] = 0.75 * image[:, lq_image_slice]
-        return image
-
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-        out['disable_time_r'] = comfy.conds.CONDConstant(False)
-        return out
-
-class Kandinsky5(BaseModel):
-    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
-        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.kandinsky5.model.Kandinsky5)
-
-    def encode_adm(self, **kwargs):
-        return kwargs["pooled_output"]
-
-    def concat_cond(self, **kwargs):
-        noise = kwargs.get("noise", None)
-        device = kwargs["device"]
-        image = torch.zeros_like(noise)
-
-        mask = kwargs.get("concat_mask", kwargs.get("denoise_mask", None))
-        if mask is None:
-            mask = torch.zeros_like(noise)[:, :1]
-        else:
-            mask = 1.0 - mask
-            mask = utils.common_upscale(mask.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
-            if mask.shape[-3] < noise.shape[-3]:
-                mask = torch.nn.functional.pad(mask, (0, 0, 0, 0, 0, noise.shape[-3] - mask.shape[-3]), mode='constant', value=0)
-            mask = utils.resize_to_batch_size(mask, noise.shape[0])
-
-        return torch.cat((image, mask), dim=1)
-
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-        attention_mask = kwargs.get("attention_mask", None)
-        if attention_mask is not None:
-            out['attention_mask'] = comfy.conds.CONDRegular(attention_mask)
-        cross_attn = kwargs.get("cross_attn", None)
-        if cross_attn is not None:
-            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
-
-        time_dim_replace = kwargs.get("time_dim_replace", None)
-        if time_dim_replace is not None:
-            out['time_dim_replace'] = comfy.conds.CONDRegular(self.process_latent_in(time_dim_replace))
-
-        return out
-
-class Kandinsky5Image(Kandinsky5):
-    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
-        super().__init__(model_config, model_type, device=device)
-
-    def concat_cond(self, **kwargs):
-        return None

comfy/model_detection.py

@@ -6,6 +6,20 @@ import math
 import logging
 import torch
 
+
+def detect_layer_quantization(metadata):
+    quant_key = "_quantization_metadata"
+    if metadata is not None and quant_key in metadata:
+        quant_metadata = metadata.pop(quant_key)
+        quant_metadata = json.loads(quant_metadata)
+        if isinstance(quant_metadata, dict) and "layers" in quant_metadata:
+            logging.info(f"Found quantization metadata (version {quant_metadata.get('format_version', 'unknown')})")
+            return quant_metadata["layers"]
+        else:
+            raise ValueError("Invalid quantization metadata format")
+    return None
+
+
 def count_blocks(state_dict_keys, prefix_string):
     count = 0
     while True:
@@ -172,73 +186,30 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
 
         guidance_keys = list(filter(lambda a: a.startswith("{}guidance_in.".format(key_prefix)), state_dict_keys))
         dit_config["guidance_embed"] = len(guidance_keys) > 0
-
-        # HunyuanVideo 1.5
-        if '{}cond_type_embedding.weight'.format(key_prefix) in state_dict_keys:
-            dit_config["use_cond_type_embedding"] = True
-        else:
-            dit_config["use_cond_type_embedding"] = False
-        if '{}vision_in.proj.0.weight'.format(key_prefix) in state_dict_keys:
-            dit_config["vision_in_dim"] = state_dict['{}vision_in.proj.0.weight'.format(key_prefix)].shape[0]
-            dit_config["meanflow_sum"] = True
-        else:
-            dit_config["vision_in_dim"] = None
-            dit_config["meanflow_sum"] = False
         return dit_config
 
     if '{}double_blocks.0.img_attn.norm.key_norm.scale'.format(key_prefix) in state_dict_keys and ('{}img_in.weight'.format(key_prefix) in state_dict_keys or f"{key_prefix}distilled_guidance_layer.norms.0.scale" in state_dict_keys): #Flux, Chroma or Chroma Radiance (has no img_in.weight)
         dit_config = {}
-        if '{}double_stream_modulation_img.lin.weight'.format(key_prefix) in state_dict_keys:
-            dit_config["image_model"] = "flux2"
-            dit_config["axes_dim"] = [32, 32, 32, 32]
-            dit_config["num_heads"] = 48
-            dit_config["mlp_ratio"] = 3.0
-            dit_config["theta"] = 2000
-            dit_config["out_channels"] = 128
-            dit_config["global_modulation"] = True
-            dit_config["mlp_silu_act"] = True
-            dit_config["qkv_bias"] = False
-            dit_config["ops_bias"] = False
-            dit_config["default_ref_method"] = "index"
-            dit_config["ref_index_scale"] = 10.0
-            dit_config["txt_ids_dims"] = [3]
-            patch_size = 1
-        else:
-            dit_config["image_model"] = "flux"
-            dit_config["axes_dim"] = [16, 56, 56]
-            dit_config["num_heads"] = 24
-            dit_config["mlp_ratio"] = 4.0
-            dit_config["theta"] = 10000
-            dit_config["out_channels"] = 16
-            dit_config["qkv_bias"] = True
-            dit_config["txt_ids_dims"] = []
-            patch_size = 2
-
+        dit_config["image_model"] = "flux"
         dit_config["in_channels"] = 16
-        dit_config["hidden_size"] = 3072
-        dit_config["context_in_dim"] = 4096
-
+        patch_size = 2
         dit_config["patch_size"] = patch_size
         in_key = "{}img_in.weight".format(key_prefix)
         if in_key in state_dict_keys:
-            w = state_dict[in_key]
-            dit_config["in_channels"] = w.shape[1] // (patch_size * patch_size)
-            dit_config["hidden_size"] = w.shape[0]
-
-        txt_in_key = "{}txt_in.weight".format(key_prefix)
-        if txt_in_key in state_dict_keys:
-            w = state_dict[txt_in_key]
-            dit_config["context_in_dim"] = w.shape[1]
-            dit_config["hidden_size"] = w.shape[0]
-
+            dit_config["in_channels"] = state_dict[in_key].shape[1] // (patch_size * patch_size)
+        dit_config["out_channels"] = 16
         vec_in_key = '{}vector_in.in_layer.weight'.format(key_prefix)
         if vec_in_key in state_dict_keys:
             dit_config["vec_in_dim"] = state_dict[vec_in_key].shape[1]
-        else:
-            dit_config["vec_in_dim"] = None
-
+        dit_config["context_in_dim"] = 4096
+        dit_config["hidden_size"] = 3072
+        dit_config["mlp_ratio"] = 4.0
+        dit_config["num_heads"] = 24
         dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')
         dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')
+        dit_config["axes_dim"] = [16, 56, 56]
+        dit_config["theta"] = 10000
+        dit_config["qkv_bias"] = True
         if '{}distilled_guidance_layer.0.norms.0.scale'.format(key_prefix) in state_dict_keys or '{}distilled_guidance_layer.norms.0.scale'.format(key_prefix) in state_dict_keys: #Chroma
             dit_config["image_model"] = "chroma"
             dit_config["in_channels"] = 64
@@ -259,17 +230,8 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
                 dit_config["nerf_tile_size"] = 512
                 dit_config["nerf_final_head_type"] = "conv" if f"{key_prefix}nerf_final_layer_conv.norm.scale" in state_dict_keys else "linear"
                 dit_config["nerf_embedder_dtype"] = torch.float32
-                if "{}__x0__".format(key_prefix) in state_dict_keys: # x0 pred
-                    dit_config["use_x0"] = True
-                else:
-                    dit_config["use_x0"] = False
         else:
             dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys
-            dit_config["yak_mlp"] = '{}double_blocks.0.img_mlp.gate_proj.weight'.format(key_prefix) in state_dict_keys
-            dit_config["txt_norm"] = "{}txt_norm.scale".format(key_prefix) in state_dict_keys
-            if dit_config["yak_mlp"] and dit_config["txt_norm"]:  # Ovis model
-                dit_config["txt_ids_dims"] = [1, 2]
-
         return dit_config
 
     if '{}t5_yproj.weight'.format(key_prefix) in state_dict_keys: #Genmo mochi preview
@@ -416,35 +378,14 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
         dit_config["image_model"] = "lumina2"
         dit_config["patch_size"] = 2
         dit_config["in_channels"] = 16
-        w = state_dict['{}cap_embedder.1.weight'.format(key_prefix)]
-        dit_config["dim"] = w.shape[0]
-        dit_config["cap_feat_dim"] = w.shape[1]
+        dit_config["dim"] = 2304
+        dit_config["cap_feat_dim"] = state_dict['{}cap_embedder.1.weight'.format(key_prefix)].shape[1]
         dit_config["n_layers"] = count_blocks(state_dict_keys, '{}layers.'.format(key_prefix) + '{}.')
+        dit_config["n_heads"] = 24
+        dit_config["n_kv_heads"] = 8
         dit_config["qk_norm"] = True
-
-        if dit_config["dim"] == 2304: # Original Lumina 2
-            dit_config["n_heads"] = 24
-            dit_config["n_kv_heads"] = 8
-            dit_config["axes_dims"] = [32, 32, 32]
-            dit_config["axes_lens"] = [300, 512, 512]
-            dit_config["rope_theta"] = 10000.0
-            dit_config["ffn_dim_multiplier"] = 4.0
-            ctd_weight = state_dict.get('{}clip_text_pooled_proj.0.weight'.format(key_prefix), None)
-            if ctd_weight is not None:  # NewBie
-                dit_config["clip_text_dim"] = ctd_weight.shape[0]
-                # NewBie also sets axes_lens = [1024, 512, 512] but it's not used in ComfyUI
-        elif dit_config["dim"] == 3840:  # Z image
-            dit_config["n_heads"] = 30
-            dit_config["n_kv_heads"] = 30
-            dit_config["axes_dims"] = [32, 48, 48]
-            dit_config["axes_lens"] = [1536, 512, 512]
-            dit_config["rope_theta"] = 256.0
-            dit_config["ffn_dim_multiplier"] = (8.0 / 3.0)
-            dit_config["z_image_modulation"] = True
-            dit_config["time_scale"] = 1000.0
-            if '{}cap_pad_token'.format(key_prefix) in state_dict_keys:
-                dit_config["pad_tokens_multiple"] = 32
-
+        dit_config["axes_dims"] = [32, 32, 32]
+        dit_config["axes_lens"] = [300, 512, 512]
         return dit_config
 
     if '{}head.modulation'.format(key_prefix) in state_dict_keys:  # Wan 2.1
@@ -619,29 +560,6 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
         dit_config["image_model"] = "qwen_image"
         dit_config["in_channels"] = state_dict['{}img_in.weight'.format(key_prefix)].shape[1]
         dit_config["num_layers"] = count_blocks(state_dict_keys, '{}transformer_blocks.'.format(key_prefix) + '{}.')
-        if "{}__index_timestep_zero__".format(key_prefix) in state_dict_keys:  # 2511
-            dit_config["default_ref_method"] = "index_timestep_zero"
-        if "{}time_text_embed.addition_t_embedding.weight".format(key_prefix) in state_dict_keys:  # Layered
-            dit_config["use_additional_t_cond"] = True
-            dit_config["default_ref_method"] = "negative_index"
-        return dit_config
-
-    if '{}visual_transformer_blocks.0.cross_attention.key_norm.weight'.format(key_prefix) in state_dict_keys: # Kandinsky 5
-        dit_config = {}
-        model_dim = state_dict['{}visual_embeddings.in_layer.bias'.format(key_prefix)].shape[0]
-        dit_config["model_dim"] = model_dim
-        if model_dim in [4096, 2560]: # pro video and lite image
-            dit_config["axes_dims"] = (32, 48, 48)
-            if model_dim == 2560: # lite image
-                dit_config["rope_scale_factor"] = (1.0, 1.0, 1.0)
-        elif model_dim == 1792: # lite video
-            dit_config["axes_dims"] = (16, 24, 24)
-        dit_config["time_dim"] = state_dict['{}time_embeddings.in_layer.bias'.format(key_prefix)].shape[0]
-        dit_config["image_model"] = "kandinsky5"
-        dit_config["ff_dim"] = state_dict['{}visual_transformer_blocks.0.feed_forward.in_layer.weight'.format(key_prefix)].shape[0]
-        dit_config["visual_embed_dim"] = state_dict['{}visual_embeddings.in_layer.weight'.format(key_prefix)].shape[1]
-        dit_config["num_text_blocks"] = count_blocks(state_dict_keys, '{}text_transformer_blocks.'.format(key_prefix) + '{}.')
-        dit_config["num_visual_blocks"] = count_blocks(state_dict_keys, '{}visual_transformer_blocks.'.format(key_prefix) + '{}.')
         return dit_config
 
     if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
@@ -786,11 +704,22 @@ def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=Fal
     if model_config is None and use_base_if_no_match:
         model_config = comfy.supported_models_base.BASE(unet_config)
 
+    scaled_fp8_key = "{}scaled_fp8".format(unet_key_prefix)
+    if scaled_fp8_key in state_dict:
+        scaled_fp8_weight = state_dict.pop(scaled_fp8_key)
+        model_config.scaled_fp8 = scaled_fp8_weight.dtype
+        if model_config.scaled_fp8 == torch.float32:
+            model_config.scaled_fp8 = torch.float8_e4m3fn
+        if scaled_fp8_weight.nelement() == 2:
+            model_config.optimizations["fp8"] = False
+        else:
+            model_config.optimizations["fp8"] = True
+
     # Detect per-layer quantization (mixed precision)
-    quant_config = comfy.utils.detect_layer_quantization(state_dict, unet_key_prefix)
-    if quant_config:
-        model_config.quant_config = quant_config
-        logging.info("Detected mixed precision quantization")
+    layer_quant_config = detect_layer_quantization(metadata)
+    if layer_quant_config:
+        model_config.layer_quant_config = layer_quant_config
+        logging.info(f"Detected mixed precision quantization: {len(layer_quant_config)} layers quantized")
 
     return model_config
 

comfy/model_management.py

@@ -26,7 +26,6 @@ import importlib
 import platform
 import weakref
 import gc
-import os
 
 class VRAMState(Enum):
     DISABLED = 0    #No vram present: no need to move models to vram
@@ -334,15 +333,13 @@ except:
 SUPPORT_FP8_OPS = args.supports_fp8_compute
 
 AMD_RDNA2_AND_OLDER_ARCH = ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]
-AMD_ENABLE_MIOPEN_ENV = 'COMFYUI_ENABLE_MIOPEN'
 
 try:
     if is_amd():
         arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName
         if not (any((a in arch) for a in AMD_RDNA2_AND_OLDER_ARCH)):
-            if os.getenv(AMD_ENABLE_MIOPEN_ENV) != '1':
-                torch.backends.cudnn.enabled = False  # Seems to improve things a lot on AMD
-                logging.info("Set: torch.backends.cudnn.enabled = False for better AMD performance.")
+            torch.backends.cudnn.enabled = False  # Seems to improve things a lot on AMD
+            logging.info("Set: torch.backends.cudnn.enabled = False for better AMD performance.")
 
         try:
             rocm_version = tuple(map(int, str(torch.version.hip).split(".")[:2]))
@@ -507,7 +504,6 @@ class LoadedModel:
         if use_more_vram == 0:
             use_more_vram = 1e32
         self.model_use_more_vram(use_more_vram, force_patch_weights=force_patch_weights)
-
         real_model = self.model.model
 
         if is_intel_xpu() and not args.disable_ipex_optimize and 'ipex' in globals() and real_model is not None:
@@ -692,11 +688,8 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
             loaded_memory = loaded_model.model_loaded_memory()
             current_free_mem = get_free_memory(torch_dev) + loaded_memory
 
-            lowvram_model_memory = max(0, (current_free_mem - minimum_memory_required), min(current_free_mem * MIN_WEIGHT_MEMORY_RATIO, current_free_mem - minimum_inference_memory()))
-            lowvram_model_memory = lowvram_model_memory - loaded_memory
-
-            if lowvram_model_memory == 0:
-                lowvram_model_memory = 0.1
+            lowvram_model_memory = max(128 * 1024 * 1024, (current_free_mem - minimum_memory_required), min(current_free_mem * MIN_WEIGHT_MEMORY_RATIO, current_free_mem - minimum_inference_memory()))
+            lowvram_model_memory = max(0.1, lowvram_model_memory - loaded_memory)
 
         if vram_set_state == VRAMState.NO_VRAM:
             lowvram_model_memory = 0.1
@@ -1015,18 +1008,9 @@ def force_channels_last():
 
 
 STREAMS = {}
-NUM_STREAMS = 0
-if args.async_offload is not None:
-    NUM_STREAMS = args.async_offload
-else:
-    #  Enable by default on Nvidia
-    if is_nvidia():
-        NUM_STREAMS = 2
-
-if args.disable_async_offload:
-    NUM_STREAMS = 0
-
-if NUM_STREAMS > 0:
+NUM_STREAMS = 1
+if args.async_offload:
+    NUM_STREAMS = 2
     logging.info("Using async weight offloading with {} streams".format(NUM_STREAMS))
 
 def current_stream(device):
@@ -1042,10 +1026,7 @@ def current_stream(device):
 stream_counters = {}
 def get_offload_stream(device):
     stream_counter = stream_counters.get(device, 0)
-    if NUM_STREAMS == 0:
-        return None
-
-    if torch.compiler.is_compiling():
+    if NUM_STREAMS <= 1:
         return None
 
     if device in STREAMS:
@@ -1058,9 +1039,7 @@ def get_offload_stream(device):
     elif is_device_cuda(device):
         ss = []
         for k in range(NUM_STREAMS):
-            s1 = torch.cuda.Stream(device=device, priority=0)
-            s1.as_context = torch.cuda.stream
-            ss.append(s1)
+            ss.append(torch.cuda.Stream(device=device, priority=0))
         STREAMS[device] = ss
         s = ss[stream_counter]
         stream_counters[device] = stream_counter
@@ -1068,9 +1047,7 @@ def get_offload_stream(device):
     elif is_device_xpu(device):
         ss = []
         for k in range(NUM_STREAMS):
-            s1 = torch.xpu.Stream(device=device, priority=0)
-            s1.as_context = torch.xpu.stream
-            ss.append(s1)
+            ss.append(torch.xpu.Stream(device=device, priority=0))
         STREAMS[device] = ss
         s = ss[stream_counter]
         stream_counters[device] = stream_counter
@@ -1088,19 +1065,12 @@ def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, str
             if dtype is None or weight.dtype == dtype:
                 return weight
         if stream is not None:
-            wf_context = stream
-            if hasattr(wf_context, "as_context"):
-                wf_context = wf_context.as_context(stream)
-            with wf_context:
+            with stream:
                 return weight.to(dtype=dtype, copy=copy)
         return weight.to(dtype=dtype, copy=copy)
 
-
     if stream is not None:
-        wf_context = stream
-        if hasattr(wf_context, "as_context"):
-            wf_context = wf_context.as_context(stream)
-        with wf_context:
+        with stream:
             r = torch.empty_like(weight, dtype=dtype, device=device)
             r.copy_(weight, non_blocking=non_blocking)
     else:
@@ -1124,16 +1094,12 @@ if not args.disable_pinned_memory:
             MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.95
         logging.info("Enabled pinned memory {}".format(MAX_PINNED_MEMORY // (1024 * 1024)))
 
-PINNING_ALLOWED_TYPES = set(["Parameter", "QuantizedTensor"])
 
 def pin_memory(tensor):
     global TOTAL_PINNED_MEMORY
     if MAX_PINNED_MEMORY <= 0:
         return False
 
-    if type(tensor).__name__ not in PINNING_ALLOWED_TYPES:
-        return False
-
     if not is_device_cpu(tensor.device):
         return False
 
@@ -1143,17 +1109,11 @@ def pin_memory(tensor):
         #on the GPU async. So dont trust the CUDA API and guard here
         return False
 
-    if not tensor.is_contiguous():
-        return False
-
     size = tensor.numel() * tensor.element_size()
     if (TOTAL_PINNED_MEMORY + size) > MAX_PINNED_MEMORY:
         return False
 
     ptr = tensor.data_ptr()
-    if ptr == 0:
-        return False
-
     if torch.cuda.cudart().cudaHostRegister(ptr, size, 1) == 0:
         PINNED_MEMORY[ptr] = size
         TOTAL_PINNED_MEMORY += size
@@ -1495,20 +1455,6 @@ def extended_fp16_support():
 
     return True
 
-LORA_COMPUTE_DTYPES = {}
-def lora_compute_dtype(device):
-    dtype = LORA_COMPUTE_DTYPES.get(device, None)
-    if dtype is not None:
-        return dtype
-
-    if should_use_fp16(device):
-        dtype = torch.float16
-    else:
-        dtype = torch.float32
-
-    LORA_COMPUTE_DTYPES[device] = dtype
-    return dtype
-
 def soft_empty_cache(force=False):
     global cpu_state
     if cpu_state == CPUState.MPS:

comfy/model_patcher.py

@@ -35,7 +35,6 @@ import comfy.model_management
 import comfy.patcher_extension
 import comfy.utils
 from comfy.comfy_types import UnetWrapperFunction
-from comfy.quant_ops import QuantizedTensor
 from comfy.patcher_extension import CallbacksMP, PatcherInjection, WrappersMP
 
 
@@ -127,23 +126,27 @@ class LowVramPatch:
     def __init__(self, key, patches, convert_func=None, set_func=None):
         self.key = key
         self.patches = patches
-        self.convert_func = convert_func # TODO: remove
+        self.convert_func = convert_func
         self.set_func = set_func
 
     def __call__(self, weight):
-        return comfy.lora.calculate_weight(self.patches[self.key], weight, self.key, intermediate_dtype=weight.dtype)
+        intermediate_dtype = weight.dtype
+        if self.convert_func is not None:
+            weight = self.convert_func(weight.to(dtype=torch.float32, copy=True), inplace=True)
 
-LOWVRAM_PATCH_ESTIMATE_MATH_FACTOR = 2
+        if intermediate_dtype not in [torch.float32, torch.float16, torch.bfloat16]: #intermediate_dtype has to be one that is supported in math ops
+            intermediate_dtype = torch.float32
+            out = comfy.lora.calculate_weight(self.patches[self.key], weight.to(intermediate_dtype), self.key, intermediate_dtype=intermediate_dtype)
+            if self.set_func is None:
+                return comfy.float.stochastic_rounding(out, weight.dtype, seed=string_to_seed(self.key))
+            else:
+                return self.set_func(out, seed=string_to_seed(self.key), return_weight=True)
 
-def low_vram_patch_estimate_vram(model, key):
-    weight, set_func, convert_func = get_key_weight(model, key)
-    if weight is None:
-        return 0
-    model_dtype = getattr(model, "manual_cast_dtype", torch.float32)
-    if model_dtype is None:
-        model_dtype = weight.dtype
-
-    return weight.numel() * model_dtype.itemsize * LOWVRAM_PATCH_ESTIMATE_MATH_FACTOR
+        out = comfy.lora.calculate_weight(self.patches[self.key], weight, self.key, intermediate_dtype=intermediate_dtype)
+        if self.set_func is not None:
+            return self.set_func(out, seed=string_to_seed(self.key), return_weight=True).to(dtype=intermediate_dtype)
+        else:
+            return out
 
 def get_key_weight(model, key):
     set_func = None
@@ -228,6 +231,7 @@ class ModelPatcher:
         self.object_patches_backup = {}
         self.weight_wrapper_patches = {}
         self.model_options = {"transformer_options":{}}
+        self.model_size()
         self.load_device = load_device
         self.offload_device = offload_device
         self.weight_inplace_update = weight_inplace_update
@@ -266,9 +270,6 @@ class ModelPatcher:
         if not hasattr(self.model, 'current_weight_patches_uuid'):
             self.model.current_weight_patches_uuid = None
 
-        if not hasattr(self.model, 'model_offload_buffer_memory'):
-            self.model.model_offload_buffer_memory = 0
-
     def model_size(self):
         if self.size > 0:
             return self.size
@@ -285,7 +286,7 @@ class ModelPatcher:
         return self.model.lowvram_patch_counter
 
     def clone(self):
-        n = self.__class__(self.model, self.load_device, self.offload_device, self.model_size(), weight_inplace_update=self.weight_inplace_update)
+        n = self.__class__(self.model, self.load_device, self.offload_device, self.size, weight_inplace_update=self.weight_inplace_update)
         n.patches = {}
         for k in self.patches:
             n.patches[k] = self.patches[k][:]
@@ -454,9 +455,6 @@ class ModelPatcher:
     def set_model_post_input_patch(self, patch):
         self.set_model_patch(patch, "post_input")
 
-    def set_model_noise_refiner_patch(self, patch):
-        self.set_model_patch(patch, "noise_refiner")
-
     def set_model_rope_options(self, scale_x, shift_x, scale_y, shift_y, scale_t, shift_t, **kwargs):
         rope_options = self.model_options["transformer_options"].get("rope_options", {})
         rope_options["scale_x"] = scale_x
@@ -621,11 +619,10 @@ class ModelPatcher:
         if key not in self.backup:
             self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update)
 
-        temp_dtype = comfy.model_management.lora_compute_dtype(device_to)
         if device_to is not None:
-            temp_weight = comfy.model_management.cast_to_device(weight, device_to, temp_dtype, copy=True)
+            temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
         else:
-            temp_weight = weight.to(temp_dtype, copy=True)
+            temp_weight = weight.to(torch.float32, copy=True)
         if convert_func is not None:
             temp_weight = convert_func(temp_weight, inplace=True)
 
@@ -666,22 +663,7 @@ class ModelPatcher:
                     skip = True # skip random weights in non leaf modules
                     break
             if not skip and (hasattr(m, "comfy_cast_weights") or len(params) > 0):
-                module_mem = comfy.model_management.module_size(m)
-                module_offload_mem = module_mem
-                if hasattr(m, "comfy_cast_weights"):
-                    def check_module_offload_mem(key):
-                        if key in self.patches:
-                            return low_vram_patch_estimate_vram(self.model, key)
-                        model_dtype = getattr(self.model, "manual_cast_dtype", None)
-                        weight, _, _ = get_key_weight(self.model, key)
-                        if model_dtype is None or weight is None:
-                            return 0
-                        if (weight.dtype != model_dtype or isinstance(weight, QuantizedTensor)):
-                            return weight.numel() * model_dtype.itemsize
-                        return 0
-                    module_offload_mem += check_module_offload_mem("{}.weight".format(n))
-                    module_offload_mem += check_module_offload_mem("{}.bias".format(n))
-                loading.append((module_offload_mem, module_mem, n, m, params))
+                loading.append((comfy.model_management.module_size(m), n, m, params))
         return loading
 
     def load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False):
@@ -695,22 +677,20 @@ class ModelPatcher:
 
             load_completely = []
             offloaded = []
-            offload_buffer = 0
             loading.sort(reverse=True)
-            for i, x in enumerate(loading):
-                module_offload_mem, module_mem, n, m, params = x
+            for x in loading:
+                n = x[1]
+                m = x[2]
+                params = x[3]
+                module_mem = x[0]
 
                 lowvram_weight = False
 
-                potential_offload = max(offload_buffer, module_offload_mem + sum([ x1[1] for x1 in loading[i+1:i+1+comfy.model_management.NUM_STREAMS]]))
-                lowvram_fits = mem_counter + module_mem + potential_offload < lowvram_model_memory
-
                 weight_key = "{}.weight".format(n)
                 bias_key = "{}.bias".format(n)
 
                 if not full_load and hasattr(m, "comfy_cast_weights"):
-                    if not lowvram_fits:
-                        offload_buffer = potential_offload
+                    if mem_counter + module_mem >= lowvram_model_memory:
                         lowvram_weight = True
                         lowvram_counter += 1
                         lowvram_mem_counter += module_mem
@@ -744,11 +724,9 @@ class ModelPatcher:
                     if hasattr(m, "comfy_cast_weights"):
                         wipe_lowvram_weight(m)
 
-                    if full_load or lowvram_fits:
+                    if full_load or mem_counter + module_mem < lowvram_model_memory:
                         mem_counter += module_mem
                         load_completely.append((module_mem, n, m, params))
-                    else:
-                        offload_buffer = potential_offload
 
                 if cast_weight and hasattr(m, "comfy_cast_weights"):
                     m.prev_comfy_cast_weights = m.comfy_cast_weights
@@ -775,8 +753,6 @@ class ModelPatcher:
                     key = "{}.{}".format(n, param)
                     self.unpin_weight(key)
                     self.patch_weight_to_device(key, device_to=device_to)
-                if comfy.model_management.is_device_cuda(device_to):
-                    torch.cuda.synchronize()
 
                 logging.debug("lowvram: loaded module regularly {} {}".format(n, m))
                 m.comfy_patched_weights = True
@@ -791,7 +767,7 @@ class ModelPatcher:
                     self.pin_weight_to_device("{}.{}".format(n, param))
 
             if lowvram_counter > 0:
-                logging.info("loaded partially; {:.2f} MB usable, {:.2f} MB loaded, {:.2f} MB offloaded, {:.2f} MB buffer reserved, lowvram patches: {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), lowvram_mem_counter / (1024 * 1024), offload_buffer / (1024 * 1024), patch_counter))
+                logging.info("loaded partially; {:.2f} MB usable, {:.2f} MB loaded, {:.2f} MB offloaded, lowvram patches: {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), lowvram_mem_counter / (1024 * 1024), patch_counter))
                 self.model.model_lowvram = True
             else:
                 logging.info("loaded completely; {:.2f} MB usable, {:.2f} MB loaded, full load: {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), full_load))
@@ -803,7 +779,6 @@ class ModelPatcher:
             self.model.lowvram_patch_counter += patch_counter
             self.model.device = device_to
             self.model.model_loaded_weight_memory = mem_counter
-            self.model.model_offload_buffer_memory = offload_buffer
             self.model.current_weight_patches_uuid = self.patches_uuid
 
             for callback in self.get_all_callbacks(CallbacksMP.ON_LOAD):
@@ -857,7 +832,6 @@ class ModelPatcher:
                 self.model.to(device_to)
                 self.model.device = device_to
             self.model.model_loaded_weight_memory = 0
-            self.model.model_offload_buffer_memory = 0
 
             for m in self.model.modules():
                 if hasattr(m, "comfy_patched_weights"):
@@ -869,25 +843,20 @@ class ModelPatcher:
 
         self.object_patches_backup.clear()
 
-    def partially_unload(self, device_to, memory_to_free=0, force_patch_weights=False):
+    def partially_unload(self, device_to, memory_to_free=0):
         with self.use_ejected():
             hooks_unpatched = False
             memory_freed = 0
             patch_counter = 0
             unload_list = self._load_list()
             unload_list.sort()
-
-            offload_buffer = self.model.model_offload_buffer_memory
-            if len(unload_list) > 0:
-                NS = comfy.model_management.NUM_STREAMS
-                offload_weight_factor = [ min(offload_buffer / (NS + 1), unload_list[0][1]) ] * NS
-
             for unload in unload_list:
-                if memory_to_free + offload_buffer - self.model.model_offload_buffer_memory < memory_freed:
+                if memory_to_free < memory_freed:
                     break
-                module_offload_mem, module_mem, n, m, params = unload
-
-                potential_offload = module_offload_mem + sum(offload_weight_factor)
+                module_mem = unload[0]
+                n = unload[1]
+                m = unload[2]
+                params = unload[3]
 
                 lowvram_possible = hasattr(m, "comfy_cast_weights")
                 if hasattr(m, "comfy_patched_weights") and m.comfy_patched_weights == True:
@@ -918,40 +887,28 @@ class ModelPatcher:
                         module_mem += move_weight_functions(m, device_to)
                         if lowvram_possible:
                             if weight_key in self.patches:
-                                if force_patch_weights:
-                                    self.patch_weight_to_device(weight_key)
-                                else:
-                                    _, set_func, convert_func = get_key_weight(self.model, weight_key)
-                                    m.weight_function.append(LowVramPatch(weight_key, self.patches, convert_func, set_func))
-                                    patch_counter += 1
+                                _, set_func, convert_func = get_key_weight(self.model, weight_key)
+                                m.weight_function.append(LowVramPatch(weight_key, self.patches, convert_func, set_func))
+                                patch_counter += 1
                             if bias_key in self.patches:
-                                if force_patch_weights:
-                                    self.patch_weight_to_device(bias_key)
-                                else:
-                                    _, set_func, convert_func = get_key_weight(self.model, bias_key)
-                                    m.bias_function.append(LowVramPatch(bias_key, self.patches, convert_func, set_func))
-                                    patch_counter += 1
+                                _, set_func, convert_func = get_key_weight(self.model, bias_key)
+                                m.bias_function.append(LowVramPatch(bias_key, self.patches, convert_func, set_func))
+                                patch_counter += 1
                             cast_weight = True
 
-                        if cast_weight and hasattr(m, "comfy_cast_weights"):
+                        if cast_weight:
                             m.prev_comfy_cast_weights = m.comfy_cast_weights
                             m.comfy_cast_weights = True
                         m.comfy_patched_weights = False
                         memory_freed += module_mem
-                        offload_buffer = max(offload_buffer, potential_offload)
-                        offload_weight_factor.append(module_mem)
-                        offload_weight_factor.pop(0)
                         logging.debug("freed {}".format(n))
 
                         for param in params:
                             self.pin_weight_to_device("{}.{}".format(n, param))
 
-
             self.model.model_lowvram = True
             self.model.lowvram_patch_counter += patch_counter
             self.model.model_loaded_weight_memory -= memory_freed
-            self.model.model_offload_buffer_memory = offload_buffer
-            logging.info("Unloaded partially: {:.2f} MB freed, {:.2f} MB remains loaded, {:.2f} MB buffer reserved, lowvram patches: {}".format(memory_freed / (1024 * 1024), self.model.model_loaded_weight_memory / (1024 * 1024), offload_buffer / (1024 * 1024), self.model.lowvram_patch_counter))
             return memory_freed
 
     def partially_load(self, device_to, extra_memory=0, force_patch_weights=False):
@@ -964,9 +921,6 @@ class ModelPatcher:
                 extra_memory += (used - self.model.model_loaded_weight_memory)
 
             self.patch_model(load_weights=False)
-            if extra_memory < 0 and not unpatch_weights:
-                self.partially_unload(self.offload_device, -extra_memory, force_patch_weights=force_patch_weights)
-                return 0
             full_load = False
             if self.model.model_lowvram == False and self.model.model_loaded_weight_memory > 0:
                 self.apply_hooks(self.forced_hooks, force_apply=True)

comfy/ops.py

@@ -22,7 +22,7 @@ import comfy.model_management
 from comfy.cli_args import args, PerformanceFeature
 import comfy.float
 import comfy.rmsnorm
-import json
+import contextlib
 
 def run_every_op():
     if torch.compiler.is_compiling():
@@ -58,8 +58,7 @@ except (ModuleNotFoundError, TypeError):
 NVIDIA_MEMORY_CONV_BUG_WORKAROUND = False
 try:
     if comfy.model_management.is_nvidia():
-        cudnn_version = torch.backends.cudnn.version()
-        if (cudnn_version >= 91002 and cudnn_version < 91500) and comfy.model_management.torch_version_numeric >= (2, 9) and comfy.model_management.torch_version_numeric <= (2, 10):
+        if torch.backends.cudnn.version() >= 91002 and comfy.model_management.torch_version_numeric >= (2, 9) and comfy.model_management.torch_version_numeric <= (2, 10):
             #TODO: change upper bound version once it's fixed'
             NVIDIA_MEMORY_CONV_BUG_WORKAROUND = True
             logging.info("working around nvidia conv3d memory bug.")
@@ -78,10 +77,7 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
     # will add async-offload support to your cast and improve performance.
     if input is not None:
         if dtype is None:
-            if isinstance(input, QuantizedTensor):
-                dtype = input._layout_params["orig_dtype"]
-            else:
-                dtype = input.dtype
+            dtype = input.dtype
         if bias_dtype is None:
             bias_dtype = dtype
         if device is None:
@@ -93,6 +89,11 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
     else:
         offload_stream = None
 
+    if offload_stream is not None:
+        wf_context = offload_stream
+    else:
+        wf_context = contextlib.nullcontext()
+
     non_blocking = comfy.model_management.device_supports_non_blocking(device)
 
     weight_has_function = len(s.weight_function) > 0
@@ -104,24 +105,20 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
     if s.bias is not None:
         bias = comfy.model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=bias_has_function, stream=offload_stream)
 
+        if bias_has_function:
+            with wf_context:
+                for f in s.bias_function:
+                    bias = f(bias)
+
+    weight = weight.to(dtype=dtype)
+    if weight_has_function:
+        with wf_context:
+            for f in s.weight_function:
+                weight = f(weight)
+
     comfy.model_management.sync_stream(device, offload_stream)
-
-    bias_a = bias
-    weight_a = weight
-
-    if s.bias is not None:
-        for f in s.bias_function:
-            bias = f(bias)
-
-    if weight_has_function or weight.dtype != dtype:
-        weight = weight.to(dtype=dtype)
-        if isinstance(weight, QuantizedTensor):
-            weight = weight.dequantize()
-        for f in s.weight_function:
-            weight = f(weight)
-
     if offloadable:
-        return weight, bias, (offload_stream, weight_a, bias_a)
+        return weight, bias, offload_stream
     else:
         #Legacy function signature
         return weight, bias
@@ -130,16 +127,13 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
 def uncast_bias_weight(s, weight, bias, offload_stream):
     if offload_stream is None:
         return
-    os, weight_a, bias_a = offload_stream
-    if os is None:
-        return
-    if weight_a is not None:
-        device = weight_a.device
+    if weight is not None:
+        device = weight.device
     else:
-        if bias_a is None:
+        if bias is None:
             return
-        device = bias_a.device
-    os.wait_stream(comfy.model_management.current_stream(device))
+        device = bias.device
+    offload_stream.wait_stream(comfy.model_management.current_stream(device))
 
 
 class CastWeightBiasOp:
@@ -415,12 +409,22 @@ def fp8_linear(self, input):
 
     if input.ndim == 3 or input.ndim == 2:
         w, bias, offload_stream = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True)
-        scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
 
-        scale_input = torch.ones((), device=input.device, dtype=torch.float32)
-        input = torch.clamp(input, min=-448, max=448, out=input)
-        layout_params_weight = {'scale': scale_input, 'orig_dtype': input_dtype}
-        quantized_input = QuantizedTensor(input.to(dtype).contiguous(), "TensorCoreFP8Layout", layout_params_weight)
+        scale_weight = self.scale_weight
+        scale_input = self.scale_input
+        if scale_weight is None:
+            scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
+        else:
+            scale_weight = scale_weight.to(input.device)
+
+        if scale_input is None:
+            scale_input = torch.ones((), device=input.device, dtype=torch.float32)
+            input = torch.clamp(input, min=-448, max=448, out=input)
+            layout_params_weight = {'scale': scale_input, 'orig_dtype': input_dtype}
+            quantized_input = QuantizedTensor(input.to(dtype).contiguous(), "TensorCoreFP8Layout", layout_params_weight)
+        else:
+            scale_input = scale_input.to(input.device)
+            quantized_input = QuantizedTensor.from_float(input, "TensorCoreFP8Layout", scale=scale_input, dtype=dtype)
 
         # Wrap weight in QuantizedTensor - this enables unified dispatch
         # Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
@@ -441,7 +445,7 @@ class fp8_ops(manual_cast):
             return None
 
         def forward_comfy_cast_weights(self, input):
-            if len(self.weight_function) == 0 and len(self.bias_function) == 0:
+            if not self.training:
                 try:
                     out = fp8_linear(self, input)
                     if out is not None:
@@ -454,6 +458,59 @@ class fp8_ops(manual_cast):
             uncast_bias_weight(self, weight, bias, offload_stream)
             return x
 
+def scaled_fp8_ops(fp8_matrix_mult=False, scale_input=False, override_dtype=None):
+    logging.info("Using scaled fp8: fp8 matrix mult: {}, scale input: {}".format(fp8_matrix_mult, scale_input))
+    class scaled_fp8_op(manual_cast):
+        class Linear(manual_cast.Linear):
+            def __init__(self, *args, **kwargs):
+                if override_dtype is not None:
+                    kwargs['dtype'] = override_dtype
+                super().__init__(*args, **kwargs)
+
+            def reset_parameters(self):
+                if not hasattr(self, 'scale_weight'):
+                    self.scale_weight = torch.nn.parameter.Parameter(data=torch.ones((), device=self.weight.device, dtype=torch.float32), requires_grad=False)
+
+                if not scale_input:
+                    self.scale_input = None
+
+                if not hasattr(self, 'scale_input'):
+                    self.scale_input = torch.nn.parameter.Parameter(data=torch.ones((), device=self.weight.device, dtype=torch.float32), requires_grad=False)
+                return None
+
+            def forward_comfy_cast_weights(self, input):
+                if fp8_matrix_mult:
+                    out = fp8_linear(self, input)
+                    if out is not None:
+                        return out
+
+                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+
+                if weight.numel() < input.numel(): #TODO: optimize
+                    x = torch.nn.functional.linear(input, weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype), bias)
+                else:
+                    x = torch.nn.functional.linear(input * self.scale_weight.to(device=weight.device, dtype=weight.dtype), weight, bias)
+                uncast_bias_weight(self, weight, bias, offload_stream)
+                return x
+
+            def convert_weight(self, weight, inplace=False, **kwargs):
+                if inplace:
+                    weight *= self.scale_weight.to(device=weight.device, dtype=weight.dtype)
+                    return weight
+                else:
+                    return weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype)
+
+            def set_weight(self, weight, inplace_update=False, seed=None, return_weight=False, **kwargs):
+                weight = comfy.float.stochastic_rounding(weight / self.scale_weight.to(device=weight.device, dtype=weight.dtype), self.weight.dtype, seed=seed)
+                if return_weight:
+                    return weight
+                if inplace_update:
+                    self.weight.data.copy_(weight)
+                else:
+                    self.weight = torch.nn.Parameter(weight, requires_grad=False)
+
+    return scaled_fp8_op
+
 CUBLAS_IS_AVAILABLE = False
 try:
     from cublas_ops import CublasLinear
@@ -477,182 +534,129 @@ if CUBLAS_IS_AVAILABLE:
 # ==============================================================================
 # Mixed Precision Operations
 # ==============================================================================
-from .quant_ops import QuantizedTensor, QUANT_ALGOS
+from .quant_ops import QuantizedTensor
 
+QUANT_FORMAT_MIXINS = {
+    "float8_e4m3fn": {
+        "dtype": torch.float8_e4m3fn,
+        "layout_type": "TensorCoreFP8Layout",
+        "parameters": {
+            "weight_scale": torch.nn.Parameter(torch.zeros((), dtype=torch.float32), requires_grad=False),
+            "input_scale": torch.nn.Parameter(torch.zeros((), dtype=torch.float32), requires_grad=False),
+        }
+    }
+}
 
-def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_precision_mm=False):
-    class MixedPrecisionOps(manual_cast):
-        _quant_config = quant_config
-        _compute_dtype = compute_dtype
-        _full_precision_mm = full_precision_mm
+class MixedPrecisionOps(disable_weight_init):
+    _layer_quant_config = {}
+    _compute_dtype = torch.bfloat16
 
-        class Linear(torch.nn.Module, CastWeightBiasOp):
-            def __init__(
-                self,
-                in_features: int,
-                out_features: int,
-                bias: bool = True,
-                device=None,
-                dtype=None,
-            ) -> None:
-                super().__init__()
+    class Linear(torch.nn.Module, CastWeightBiasOp):
+        def __init__(
+            self,
+            in_features: int,
+            out_features: int,
+            bias: bool = True,
+            device=None,
+            dtype=None,
+        ) -> None:
+            super().__init__()
 
-                if dtype is None:
-                    dtype = MixedPrecisionOps._compute_dtype
+            self.factory_kwargs = {"device": device, "dtype": MixedPrecisionOps._compute_dtype}
+            # self.factory_kwargs = {"device": device, "dtype": dtype}
 
-                self.factory_kwargs = {"device": device, "dtype": dtype}
+            self.in_features = in_features
+            self.out_features = out_features
+            if bias:
+                self.bias = torch.nn.Parameter(torch.empty(out_features, **self.factory_kwargs))
+            else:
+                self.register_parameter("bias", None)
 
-                self.in_features = in_features
-                self.out_features = out_features
-                self._has_bias = bias
+            self.tensor_class = None
 
-                self.tensor_class = None
-                self._full_precision_mm = MixedPrecisionOps._full_precision_mm
+        def reset_parameters(self):
+            return None
 
-            def reset_parameters(self):
-                return None
+        def _load_from_state_dict(self, state_dict, prefix, local_metadata,
+                                  strict, missing_keys, unexpected_keys, error_msgs):
 
-            def _load_from_state_dict(self, state_dict, prefix, local_metadata,
-                                    strict, missing_keys, unexpected_keys, error_msgs):
+            device = self.factory_kwargs["device"]
+            layer_name = prefix.rstrip('.')
+            weight_key = f"{prefix}weight"
+            weight = state_dict.pop(weight_key, None)
+            if weight is None:
+                raise ValueError(f"Missing weight for layer {layer_name}")
 
-                device = self.factory_kwargs["device"]
-                layer_name = prefix.rstrip('.')
-                weight_key = f"{prefix}weight"
-                weight = state_dict.pop(weight_key, None)
-                if weight is None:
-                    raise ValueError(f"Missing weight for layer {layer_name}")
+            manually_loaded_keys = [weight_key]
 
-                manually_loaded_keys = [weight_key]
+            if layer_name not in MixedPrecisionOps._layer_quant_config:
+                self.weight = torch.nn.Parameter(weight.to(device=device, dtype=MixedPrecisionOps._compute_dtype), requires_grad=False)
+            else:
+                quant_format = MixedPrecisionOps._layer_quant_config[layer_name].get("format", None)
+                if quant_format is None:
+                    raise ValueError(f"Unknown quantization format for layer {layer_name}")
 
-                layer_conf = state_dict.pop(f"{prefix}comfy_quant", None)
-                if layer_conf is not None:
-                    layer_conf = json.loads(layer_conf.numpy().tobytes())
+                mixin = QUANT_FORMAT_MIXINS[quant_format]
+                self.layout_type = mixin["layout_type"]
 
-                if layer_conf is None:
-                    dtype = self.factory_kwargs["dtype"]
-                    self.weight = torch.nn.Parameter(weight.to(device=device, dtype=dtype), requires_grad=False)
-                    if dtype != MixedPrecisionOps._compute_dtype:
-                        self.comfy_cast_weights = True
-                    if self._has_bias:
-                        self.bias = torch.nn.Parameter(torch.empty(self.out_features, device=device, dtype=dtype))
-                    else:
-                        self.register_parameter("bias", None)
-                else:
-                    self.quant_format = layer_conf.get("format", None)
-                    if not self._full_precision_mm:
-                        self._full_precision_mm = layer_conf.get("full_precision_matrix_mult", False)
+                scale_key = f"{prefix}weight_scale"
+                layout_params = {
+                    'scale': state_dict.pop(scale_key, None),
+                    'orig_dtype': MixedPrecisionOps._compute_dtype
+                }
+                if layout_params['scale'] is not None:
+                    manually_loaded_keys.append(scale_key)
 
-                    if self.quant_format is None:
-                        raise ValueError(f"Unknown quantization format for layer {layer_name}")
+                self.weight = torch.nn.Parameter(
+                    QuantizedTensor(weight.to(device=device, dtype=mixin["dtype"]), self.layout_type, layout_params),
+                    requires_grad=False
+                )
 
-                    qconfig = QUANT_ALGOS[self.quant_format]
-                    self.layout_type = qconfig["comfy_tensor_layout"]
-
-                    weight_scale_key = f"{prefix}weight_scale"
-                    scale = state_dict.pop(weight_scale_key, None)
-                    if scale is not None:
-                        scale = scale.to(device)
-                    layout_params = {
-                        'scale': scale,
-                        'orig_dtype': MixedPrecisionOps._compute_dtype,
-                        'block_size': qconfig.get("group_size", None),
-                    }
-
-                    if scale is not None:
-                        manually_loaded_keys.append(weight_scale_key)
-
-                    self.weight = torch.nn.Parameter(
-                        QuantizedTensor(weight.to(device=device, dtype=qconfig.get("storage_t", None)), self.layout_type, layout_params),
-                        requires_grad=False
-                    )
-
-                    if self._has_bias:
-                        self.bias = torch.nn.Parameter(torch.empty(self.out_features, device=device, dtype=MixedPrecisionOps._compute_dtype))
-                    else:
-                        self.register_parameter("bias", None)
-
-                    for param_name in qconfig["parameters"]:
-                        param_key = f"{prefix}{param_name}"
-                        _v = state_dict.pop(param_key, None)
-                        if _v is None:
-                            continue
-                        self.register_parameter(param_name, torch.nn.Parameter(_v.to(device=device), requires_grad=False))
-                        manually_loaded_keys.append(param_key)
-
-                super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
-
-                for key in manually_loaded_keys:
-                    if key in missing_keys:
-                        missing_keys.remove(key)
-
-            def state_dict(self, *args, destination=None, prefix="", **kwargs):
-                sd = super().state_dict(*args, destination=destination, prefix=prefix, **kwargs)
-                if isinstance(self.weight, QuantizedTensor):
-                    sd["{}weight_scale".format(prefix)] = self.weight._layout_params['scale']
-                    quant_conf = {"format": self.quant_format}
-                    if self._full_precision_mm:
-                        quant_conf["full_precision_matrix_mult"] = True
-                    sd["{}comfy_quant".format(prefix)] = torch.tensor(list(json.dumps(quant_conf).encode('utf-8')), dtype=torch.uint8)
-                return sd
-
-            def _forward(self, input, weight, bias):
-                return torch.nn.functional.linear(input, weight, bias)
-
-            def forward_comfy_cast_weights(self, input):
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-                x = self._forward(input, weight, bias)
-                uncast_bias_weight(self, weight, bias, offload_stream)
-                return x
-
-            def forward(self, input, *args, **kwargs):
-                run_every_op()
-
-                if self._full_precision_mm or self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
-                    return self.forward_comfy_cast_weights(input, *args, **kwargs)
-                if (getattr(self, 'layout_type', None) is not None and
-                    not isinstance(input, QuantizedTensor)):
-                    input = QuantizedTensor.from_float(input, self.layout_type, scale=getattr(self, 'input_scale', None), dtype=self.weight.dtype)
-                return self._forward(input, self.weight, self.bias)
-
-            def convert_weight(self, weight, inplace=False, **kwargs):
-                if isinstance(weight, QuantizedTensor):
-                    return weight.dequantize()
-                else:
-                    return weight
-
-            def set_weight(self, weight, inplace_update=False, seed=None, return_weight=False, **kwargs):
-                if getattr(self, 'layout_type', None) is not None:
-                    weight = QuantizedTensor.from_float(weight, self.layout_type, scale="recalculate", dtype=self.weight.dtype, stochastic_rounding=seed, inplace_ops=True)
-                else:
-                    weight = weight.to(self.weight.dtype)
-                if return_weight:
-                    return weight
-
-                assert inplace_update is False  # TODO: eventually remove the inplace_update stuff
-                self.weight = torch.nn.Parameter(weight, requires_grad=False)
-
-            def _apply(self, fn, recurse=True):  # This is to get torch.compile + moving weights to another device working
-                if recurse:
-                    for module in self.children():
-                        module._apply(fn)
-
-                for key, param in self._parameters.items():
-                    if param is None:
+                for param_name, param_value in mixin["parameters"].items():
+                    param_key = f"{prefix}{param_name}"
+                    _v = state_dict.pop(param_key, None)
+                    if _v is None:
                         continue
-                    self.register_parameter(key, torch.nn.Parameter(fn(param), requires_grad=False))
-                for key, buf in self._buffers.items():
-                    if buf is not None:
-                        self._buffers[key] = fn(buf)
-                return self
+                    setattr(self, param_name, torch.nn.Parameter(_v.to(device=device), requires_grad=False))
+                    manually_loaded_keys.append(param_key)
 
-    return MixedPrecisionOps
+            super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
 
-def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, model_config=None):
-    fp8_compute = comfy.model_management.supports_fp8_compute(load_device) # TODO: if we support more ops this needs to be more granular
+            for key in manually_loaded_keys:
+                if key in missing_keys:
+                    missing_keys.remove(key)
 
-    if model_config and hasattr(model_config, 'quant_config') and model_config.quant_config:
-        logging.info("Using mixed precision operations")
-        return mixed_precision_ops(model_config.quant_config, compute_dtype, full_precision_mm=not fp8_compute)
+        def _forward(self, input, weight, bias):
+            return torch.nn.functional.linear(input, weight, bias)
+
+        def forward_comfy_cast_weights(self, input):
+            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+            x = self._forward(input, weight, bias)
+            uncast_bias_weight(self, weight, bias, offload_stream)
+            return x
+
+        def forward(self, input, *args, **kwargs):
+            run_every_op()
+
+            if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
+                return self.forward_comfy_cast_weights(input, *args, **kwargs)
+            if (getattr(self, 'layout_type', None) is not None and
+                getattr(self, 'input_scale', None) is not None and
+                not isinstance(input, QuantizedTensor)):
+                input = QuantizedTensor.from_float(input, self.layout_type, scale=self.input_scale, fp8_dtype=self.weight.dtype)
+            return self._forward(input, self.weight, self.bias)
+
+
+def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, scaled_fp8=None, model_config=None):
+    if model_config and hasattr(model_config, 'layer_quant_config') and model_config.layer_quant_config:
+        MixedPrecisionOps._layer_quant_config = model_config.layer_quant_config
+        MixedPrecisionOps._compute_dtype = compute_dtype
+        logging.info(f"Using mixed precision operations: {len(model_config.layer_quant_config)} quantized layers")
+        return MixedPrecisionOps
+
+    fp8_compute = comfy.model_management.supports_fp8_compute(load_device)
+    if scaled_fp8 is not None:
+        return scaled_fp8_ops(fp8_matrix_mult=fp8_compute and fp8_optimizations, scale_input=fp8_optimizations, override_dtype=scaled_fp8)
 
     if (
         fp8_compute and

comfy/quant_ops.py

@@ -1,7 +1,6 @@
 import torch
 import logging
 from typing import Tuple, Dict
-import comfy.float
 
 _LAYOUT_REGISTRY = {}
 _GENERIC_UTILS = {}
@@ -75,12 +74,6 @@ def _copy_layout_params(params):
             new_params[k] = v
     return new_params
 
-def _copy_layout_params_inplace(src, dst, non_blocking=False):
-    for k, v in src.items():
-        if isinstance(v, torch.Tensor):
-            dst[k].copy_(v, non_blocking=non_blocking)
-        else:
-            dst[k] = v
 
 class QuantizedLayout:
     """
@@ -229,17 +222,6 @@ class QuantizedTensor(torch.Tensor):
         new_kwargs = dequant_arg(kwargs)
         return func(*new_args, **new_kwargs)
 
-    def data_ptr(self):
-        return self._qdata.data_ptr()
-
-    def is_pinned(self):
-        return self._qdata.is_pinned()
-
-    def is_contiguous(self, *arg, **kwargs):
-        return self._qdata.is_contiguous(*arg, **kwargs)
-
-    def storage(self):
-        return self._qdata.storage()
 
 # ==============================================================================
 # Generic Utilities (Layout-Agnostic Operations)
@@ -252,6 +234,12 @@ def _create_transformed_qtensor(qt, transform_fn):
 
 
 def _handle_device_transfer(qt, target_device, target_dtype=None, target_layout=None, op_name="to"):
+    if target_dtype is not None and target_dtype != qt.dtype:
+        logging.warning(
+            f"QuantizedTensor: dtype conversion requested to {target_dtype}, "
+            f"but not supported for quantized tensors. Ignoring dtype."
+        )
+
     if target_layout is not None and target_layout != torch.strided:
         logging.warning(
             f"QuantizedTensor: layout change requested to {target_layout}, "
@@ -271,8 +259,6 @@ def _handle_device_transfer(qt, target_device, target_dtype=None, target_layout=
             logging.debug(f"QuantizedTensor.{op_name}: Moving from {current_device} to {target_device}")
             new_q_data = qt._qdata.to(device=target_device)
             new_params = _move_layout_params_to_device(qt._layout_params, target_device)
-            if target_dtype is not None:
-                new_params["orig_dtype"] = target_dtype
             new_qt = QuantizedTensor(new_q_data, qt._layout_type, new_params)
             logging.debug(f"QuantizedTensor.{op_name}: Created new tensor on {target_device}")
             return new_qt
@@ -332,15 +318,13 @@ def generic_to_dtype_layout(func, args, kwargs):
 def generic_copy_(func, args, kwargs):
     qt_dest = args[0]
     src = args[1]
-    non_blocking = args[2] if len(args) > 2 else False
+
     if isinstance(qt_dest, QuantizedTensor):
         if isinstance(src, QuantizedTensor):
             # Copy from another quantized tensor
-            qt_dest._qdata.copy_(src._qdata, non_blocking=non_blocking)
+            qt_dest._qdata.copy_(src._qdata)
             qt_dest._layout_type = src._layout_type
-            orig_dtype = qt_dest._layout_params["orig_dtype"]
-            _copy_layout_params_inplace(src._layout_params, qt_dest._layout_params, non_blocking=non_blocking)
-            qt_dest._layout_params["orig_dtype"] = orig_dtype
+            qt_dest._layout_params = _copy_layout_params(src._layout_params)
         else:
             # Copy from regular tensor - just copy raw data
             qt_dest._qdata.copy_(src)
@@ -348,42 +332,10 @@ def generic_copy_(func, args, kwargs):
     return func(*args, **kwargs)
 
 
-@register_generic_util(torch.ops.aten.to.dtype)
-def generic_to_dtype(func, args, kwargs):
-    """Handle .to(dtype) calls - dtype conversion only."""
-    src = args[0]
-    if isinstance(src, QuantizedTensor):
-        # For dtype-only conversion, just change the orig_dtype, no real cast is needed
-        target_dtype = args[1] if len(args) > 1 else kwargs.get('dtype')
-        src._layout_params["orig_dtype"] = target_dtype
-        return src
-    return func(*args, **kwargs)
-
-
 @register_generic_util(torch.ops.aten._has_compatible_shallow_copy_type.default)
 def generic_has_compatible_shallow_copy_type(func, args, kwargs):
     return True
 
-
-@register_generic_util(torch.ops.aten.empty_like.default)
-def generic_empty_like(func, args, kwargs):
-    """Empty_like operation - creates an empty tensor with the same quantized structure."""
-    qt = args[0]
-    if isinstance(qt, QuantizedTensor):
-        # Create empty tensor with same shape and dtype as the quantized data
-        hp_dtype = kwargs.pop('dtype', qt._layout_params["orig_dtype"])
-        new_qdata = torch.empty_like(qt._qdata, **kwargs)
-
-        # Handle device transfer for layout params
-        target_device = kwargs.get('device', new_qdata.device)
-        new_params = _move_layout_params_to_device(qt._layout_params, target_device)
-
-        # Update orig_dtype if dtype is specified
-        new_params['orig_dtype'] = hp_dtype
-
-        return QuantizedTensor(new_qdata, qt._layout_type, new_params)
-    return func(*args, **kwargs)
-
 # ==============================================================================
 # FP8 Layout + Operation Handlers
 # ==============================================================================
@@ -395,57 +347,37 @@ class TensorCoreFP8Layout(QuantizedLayout):
     - orig_dtype: Original dtype before quantization (for casting back)
     """
     @classmethod
-    def quantize(cls, tensor, scale=None, dtype=torch.float8_e4m3fn, stochastic_rounding=0, inplace_ops=False):
+    def quantize(cls, tensor, scale=None, dtype=torch.float8_e4m3fn):
         orig_dtype = tensor.dtype
 
-        if isinstance(scale, str) and scale == "recalculate":
-            scale = torch.amax(tensor.abs()).to(dtype=torch.float32) / torch.finfo(dtype).max
-            if tensor.dtype not in [torch.float32, torch.bfloat16]:  # Prevent scale from being too small
-                tensor_info = torch.finfo(tensor.dtype)
-                scale = (1.0 / torch.clamp((1.0 / scale), min=tensor_info.min, max=tensor_info.max))
+        if scale is None:
+            scale = torch.amax(tensor.abs()) / torch.finfo(dtype).max
 
-        if scale is not None:
-            if not isinstance(scale, torch.Tensor):
-                scale = torch.tensor(scale)
-            scale = scale.to(device=tensor.device, dtype=torch.float32)
+        if not isinstance(scale, torch.Tensor):
+            scale = torch.tensor(scale)
+        scale = scale.to(device=tensor.device, dtype=torch.float32)
 
-            if inplace_ops:
-                tensor *= (1.0 / scale).to(tensor.dtype)
-            else:
-                tensor = tensor * (1.0 / scale).to(tensor.dtype)
-        else:
-            scale = torch.ones((), device=tensor.device, dtype=torch.float32)
-
-        if stochastic_rounding > 0:
-            tensor = comfy.float.stochastic_rounding(tensor, dtype=dtype, seed=stochastic_rounding)
-        else:
-            lp_amax = torch.finfo(dtype).max
-            torch.clamp(tensor, min=-lp_amax, max=lp_amax, out=tensor)
-            tensor = tensor.to(dtype, memory_format=torch.contiguous_format)
+        tensor_scaled = tensor * (1.0 / scale).to(tensor.dtype)
+        # TODO: uncomment this if it's actually needed because the clamp has a small performance penality'
+        # lp_amax = torch.finfo(dtype).max
+        # torch.clamp(tensor_scaled, min=-lp_amax, max=lp_amax, out=tensor_scaled)
+        qdata = tensor_scaled.to(dtype, memory_format=torch.contiguous_format)
 
         layout_params = {
             'scale': scale,
             'orig_dtype': orig_dtype
         }
-        return tensor, layout_params
+        return qdata, layout_params
 
     @staticmethod
     def dequantize(qdata, scale, orig_dtype, **kwargs):
         plain_tensor = torch.ops.aten._to_copy.default(qdata, dtype=orig_dtype)
-        plain_tensor.mul_(scale)
-        return plain_tensor
+        return plain_tensor * scale
 
     @classmethod
     def get_plain_tensors(cls, qtensor):
         return qtensor._qdata, qtensor._layout_params['scale']
 
-QUANT_ALGOS = {
-    "float8_e4m3fn": {
-        "storage_t": torch.float8_e4m3fn,
-        "parameters": {"weight_scale", "input_scale"},
-        "comfy_tensor_layout": "TensorCoreFP8Layout",
-    },
-}
 
 LAYOUTS = {
     "TensorCoreFP8Layout": TensorCoreFP8Layout,

comfy/sampler_helpers.py

@@ -122,20 +122,20 @@ def estimate_memory(model, noise_shape, conds):
     minimum_memory_required = model.model.memory_required([noise_shape[0]] + list(noise_shape[1:]), cond_shapes=cond_shapes_min)
     return memory_required, minimum_memory_required
 
-def prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None, force_full_load=False):
+def prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None):
     executor = comfy.patcher_extension.WrapperExecutor.new_executor(
         _prepare_sampling,
         comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.PREPARE_SAMPLING, model_options, is_model_options=True)
     )
-    return executor.execute(model, noise_shape, conds, model_options=model_options, force_full_load=force_full_load)
+    return executor.execute(model, noise_shape, conds, model_options=model_options)
 
-def _prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None, force_full_load=False):
+def _prepare_sampling(model: ModelPatcher, noise_shape, conds, model_options=None):
     real_model: BaseModel = None
     models, inference_memory = get_additional_models(conds, model.model_dtype())
     models += get_additional_models_from_model_options(model_options)
     models += model.get_nested_additional_models()  # TODO: does this require inference_memory update?
     memory_required, minimum_memory_required = estimate_memory(model, noise_shape, conds)
-    comfy.model_management.load_models_gpu([model] + models, memory_required=memory_required + inference_memory, minimum_memory_required=minimum_memory_required + inference_memory, force_full_load=force_full_load)
+    comfy.model_management.load_models_gpu([model] + models, memory_required=memory_required + inference_memory, minimum_memory_required=minimum_memory_required + inference_memory)
     real_model = model.model
 
     return real_model, conds, models

comfy/samplers.py

@@ -720,7 +720,7 @@ class Sampler:
         sigma = float(sigmas[0])
         return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma
 
-KSAMPLER_NAMES = ["euler", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_cfg_pp", "heun", "heunpp2", "exp_heun_2_x0", "exp_heun_2_x0_sde", "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_2s_ancestral_cfg_pp", "dpmpp_sde", "dpmpp_sde_gpu",
                   "dpmpp_2m", "dpmpp_2m_cfg_pp", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_2m_sde_heun", "dpmpp_2m_sde_heun_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
                   "ipndm", "ipndm_v", "deis", "res_multistep", "res_multistep_cfg_pp", "res_multistep_ancestral", "res_multistep_ancestral_cfg_pp",
@@ -984,6 +984,9 @@ class CFGGuider:
         self.inner_model, self.conds, self.loaded_models = comfy.sampler_helpers.prepare_sampling(self.model_patcher, noise.shape, self.conds, self.model_options)
         device = self.model_patcher.load_device
 
+        if denoise_mask is not None:
+            denoise_mask = comfy.sampler_helpers.prepare_mask(denoise_mask, noise.shape, device)
+
         noise = noise.to(device)
         latent_image = latent_image.to(device)
         sigmas = sigmas.to(device)
@@ -1010,24 +1013,6 @@ class CFGGuider:
         else:
             latent_shapes = [latent_image.shape]
 
-        if denoise_mask is not None:
-            if denoise_mask.is_nested:
-                denoise_masks = denoise_mask.unbind()
-                denoise_masks = denoise_masks[:len(latent_shapes)]
-            else:
-                denoise_masks = [denoise_mask]
-
-            for i in range(len(denoise_masks), len(latent_shapes)):
-                denoise_masks.append(torch.ones(latent_shapes[i]))
-
-            for i in range(len(denoise_masks)):
-                denoise_masks[i] = comfy.sampler_helpers.prepare_mask(denoise_masks[i], latent_shapes[i], self.model_patcher.load_device)
-
-            if len(denoise_masks) > 1:
-                denoise_mask, _ = comfy.utils.pack_latents(denoise_masks)
-            else:
-                denoise_mask = denoise_masks[0]
-
         self.conds = {}
         for k in self.original_conds:
             self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k]))

comfy/sd.py

@@ -52,11 +52,6 @@ import comfy.text_encoders.ace
 import comfy.text_encoders.omnigen2
 import comfy.text_encoders.qwen_image
 import comfy.text_encoders.hunyuan_image
-import comfy.text_encoders.z_image
-import comfy.text_encoders.ovis
-import comfy.text_encoders.kandinsky5
-import comfy.text_encoders.jina_clip_2
-import comfy.text_encoders.newbie
 
 import comfy.model_patcher
 import comfy.lora
@@ -64,8 +59,6 @@ import comfy.lora_convert
 import comfy.hooks
 import comfy.t2i_adapter.adapter
 import comfy.taesd.taesd
-import comfy.taesd.taehv
-import comfy.latent_formats
 
 import comfy.ldm.flux.redux
 
@@ -101,7 +94,7 @@ def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
 
 
 class CLIP:
-    def __init__(self, target=None, embedding_directory=None, no_init=False, tokenizer_data={}, parameters=0, state_dict=[], model_options={}):
+    def __init__(self, target=None, embedding_directory=None, no_init=False, tokenizer_data={}, parameters=0, model_options={}):
         if no_init:
             return
         params = target.params.copy()
@@ -129,32 +122,9 @@ class CLIP:
 
         self.tokenizer = tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
         self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
-        #Match torch.float32 hardcode upcast in TE implemention
-        self.patcher.set_model_compute_dtype(torch.float32)
         self.patcher.hook_mode = comfy.hooks.EnumHookMode.MinVram
         self.patcher.is_clip = True
         self.apply_hooks_to_conds = None
-        if len(state_dict) > 0:
-            if isinstance(state_dict, list):
-                for c in state_dict:
-                    m, u = self.load_sd(c)
-                    if len(m) > 0:
-                        logging.warning("clip missing: {}".format(m))
-
-                    if len(u) > 0:
-                        logging.debug("clip unexpected: {}".format(u))
-            else:
-                m, u = self.load_sd(state_dict, full_model=True)
-                if len(m) > 0:
-                    m_filter = list(filter(lambda a: ".logit_scale" not in a and ".transformer.text_projection.weight" not in a, m))
-                    if len(m_filter) > 0:
-                        logging.warning("clip missing: {}".format(m))
-                    else:
-                        logging.debug("clip missing: {}".format(m))
-
-                if len(u) > 0:
-                    logging.debug("clip unexpected {}:".format(u))
-
         if params['device'] == load_device:
             model_management.load_models_gpu([self.patcher], force_full_load=True)
         self.layer_idx = None
@@ -219,7 +189,6 @@ class CLIP:
                 self.cond_stage_model.set_clip_options({"projected_pooled": False})
 
             self.load_model()
-            self.cond_stage_model.set_clip_options({"execution_device": self.patcher.load_device})
             all_hooks.reset()
             self.patcher.patch_hooks(None)
             if show_pbar:
@@ -267,7 +236,6 @@ class CLIP:
             self.cond_stage_model.set_clip_options({"projected_pooled": False})
 
         self.load_model()
-        self.cond_stage_model.set_clip_options({"execution_device": self.patcher.load_device})
         o = self.cond_stage_model.encode_token_weights(tokens)
         cond, pooled = o[:2]
         if return_dict:
@@ -323,7 +291,6 @@ class VAE:
         self.latent_channels = 4
         self.latent_dim = 2
         self.output_channels = 3
-        self.pad_channel_value = None
         self.process_input = lambda image: image * 2.0 - 1.0
         self.process_output = lambda image: torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)
         self.working_dtypes = [torch.bfloat16, torch.float32]
@@ -389,7 +356,7 @@ class VAE:
 
                     self.memory_used_encode = lambda shape, dtype: (700 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
                     self.memory_used_decode = lambda shape, dtype: (700 * shape[2] * shape[3] * 32 * 32) * model_management.dtype_size(dtype)
-                elif sd['decoder.conv_in.weight'].shape[1] == 32 and sd['decoder.conv_in.weight'].ndim == 5:
+                elif sd['decoder.conv_in.weight'].shape[1] == 32:
                     ddconfig = {"block_out_channels": [128, 256, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 16, "ffactor_temporal": 4, "downsample_match_channel": True, "upsample_match_channel": True, "refiner_vae": False}
                     self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
                     self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
@@ -415,17 +382,6 @@ class VAE:
                         self.upscale_ratio = 4
 
                     self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
-                    if 'decoder.post_quant_conv.weight' in sd:
-                        sd = comfy.utils.state_dict_prefix_replace(sd, {"decoder.post_quant_conv.": "post_quant_conv.", "encoder.quant_conv.": "quant_conv."})
-
-                    if 'bn.running_mean' in sd:
-                        ddconfig["batch_norm_latent"] = True
-                        self.downscale_ratio *= 2
-                        self.upscale_ratio *= 2
-                        self.latent_channels *= 4
-                        old_memory_used_decode = self.memory_used_decode
-                        self.memory_used_decode = lambda shape, dtype: old_memory_used_decode(shape, dtype) *  4.0
-
                     if 'post_quant_conv.weight' in sd:
                         self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1])
                     else:
@@ -438,7 +394,6 @@ class VAE:
                 self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * 2048) * model_management.dtype_size(dtype)
                 self.latent_channels = 64
                 self.output_channels = 2
-                self.pad_channel_value = "replicate"
                 self.upscale_ratio = 2048
                 self.downscale_ratio =  2048
                 self.latent_dim = 1
@@ -486,20 +441,20 @@ class VAE:
             elif "decoder.conv_in.conv.weight" in sd and sd['decoder.conv_in.conv.weight'].shape[1] == 32:
                 ddconfig = {"block_out_channels": [128, 256, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 16, "ffactor_temporal": 4, "downsample_match_channel": True, "upsample_match_channel": True}
                 ddconfig['z_channels'] = sd["decoder.conv_in.conv.weight"].shape[1]
-                self.latent_channels = 32
+                self.latent_channels = 64
                 self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 16, 16)
                 self.upscale_index_formula = (4, 16, 16)
                 self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 16, 16)
                 self.downscale_index_formula = (4, 16, 16)
                 self.latent_dim = 3
-                self.not_video = False
+                self.not_video = True
                 self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
                 self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.EmptyRegularizer"},
                                                             encoder_config={'target': "comfy.ldm.hunyuan_video.vae_refiner.Encoder", 'params': ddconfig},
                                                             decoder_config={'target': "comfy.ldm.hunyuan_video.vae_refiner.Decoder", 'params': ddconfig})
 
-                self.memory_used_encode = lambda shape, dtype: (1400 * 9 * shape[-2] * shape[-1]) * model_management.dtype_size(dtype)
-                self.memory_used_decode = lambda shape, dtype: (3600 * 4 * shape[-2] * shape[-1] * 16 * 16) * model_management.dtype_size(dtype)
+                self.memory_used_encode = lambda shape, dtype: (1400 * shape[-2] * shape[-1]) * model_management.dtype_size(dtype)
+                self.memory_used_decode = lambda shape, dtype: (1400 * shape[-3] * shape[-2] * shape[-1] * 16 * 16) * model_management.dtype_size(dtype)
             elif "decoder.conv_in.conv.weight" in sd:
                 ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
                 ddconfig["conv3d"] = True
@@ -511,10 +466,8 @@ class VAE:
                 self.latent_dim = 3
                 self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.conv.weight"].shape[1]
                 self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1])
-                #This is likely to significantly over-estimate with single image or low frame counts as the
-                #implementation is able to completely skip caching. Rework if used as an image only VAE
-                self.memory_used_decode = lambda shape, dtype: (2800 * min(8, ((shape[2] - 1) * 4) + 1) * shape[3] * shape[4] * (8 * 8)) * model_management.dtype_size(dtype)
-                self.memory_used_encode = lambda shape, dtype: (1400 * min(9, shape[2]) * shape[3] * shape[4]) * model_management.dtype_size(dtype)
+                self.memory_used_decode = lambda shape, dtype: (1500 * shape[2] * shape[3] * shape[4] * (4 * 8 * 8)) * model_management.dtype_size(dtype)
+                self.memory_used_encode = lambda shape, dtype: (900 * max(shape[2], 2) * shape[3] * shape[4]) * model_management.dtype_size(dtype)
                 self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
             elif "decoder.unpatcher3d.wavelets" in sd:
                 self.upscale_ratio = (lambda a: max(0, a * 8 - 7), 8, 8)
@@ -543,22 +496,17 @@ class VAE:
                     self.memory_used_encode = lambda shape, dtype: 3300 * shape[3] * shape[4] * model_management.dtype_size(dtype)
                     self.memory_used_decode = lambda shape, dtype: 8000 * shape[3] * shape[4] * (16 * 16) * model_management.dtype_size(dtype)
                 else:  # Wan 2.1 VAE
-                    dim = sd["decoder.head.0.gamma"].shape[0]
                     self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
                     self.upscale_index_formula = (4, 8, 8)
                     self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)
                     self.downscale_index_formula = (4, 8, 8)
                     self.latent_dim = 3
                     self.latent_channels = 16
-                    self.output_channels = sd["encoder.conv1.weight"].shape[1]
-                    self.pad_channel_value = 1.0
-                    ddconfig = {"dim": dim, "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True], "image_channels": self.output_channels, "dropout": 0.0}
+                    ddconfig = {"dim": 96, "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True], "dropout": 0.0}
                     self.first_stage_model = comfy.ldm.wan.vae.WanVAE(**ddconfig)
                     self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
-                    self.memory_used_encode = lambda shape, dtype: (1500 if shape[2]<=4 else 6000) * shape[3] * shape[4] * model_management.dtype_size(dtype)
-                    self.memory_used_decode = lambda shape, dtype: (2200 if shape[2]<=4 else 7000) * shape[3] * shape[4] * (8*8) * model_management.dtype_size(dtype)
-
-
+                    self.memory_used_encode = lambda shape, dtype: 6000 * shape[3] * shape[4] * model_management.dtype_size(dtype)
+                    self.memory_used_decode = lambda shape, dtype: 7000 * shape[3] * shape[4] * (8 * 8) * model_management.dtype_size(dtype)
             # Hunyuan 3d v2 2.0 & 2.1
             elif "geo_decoder.cross_attn_decoder.ln_1.bias" in sd:
 
@@ -588,7 +536,6 @@ class VAE:
                 self.memory_used_decode = lambda shape, dtype: (shape[2] * shape[3] * 87000) * model_management.dtype_size(dtype)
                 self.latent_channels = 8
                 self.output_channels = 2
-                self.pad_channel_value = "replicate"
                 self.upscale_ratio = 4096
                 self.downscale_ratio = 4096
                 self.latent_dim = 2
@@ -625,35 +572,6 @@ class VAE:
                 self.process_input = lambda audio: audio
                 self.working_dtypes = [torch.float32]
                 self.crop_input = False
-            elif "decoder.22.bias" in sd: # taehv, taew and lighttae
-                self.latent_channels = sd["decoder.1.weight"].shape[1]
-                self.latent_dim = 3
-                self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 16, 16)
-                self.upscale_index_formula = (4, 16, 16)
-                self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 16, 16)
-                self.downscale_index_formula = (4, 16, 16)
-                if self.latent_channels == 48: # Wan 2.2
-                    self.first_stage_model = comfy.taesd.taehv.TAEHV(latent_channels=self.latent_channels, latent_format=None) # taehv doesn't need scaling
-                    self.process_input = lambda image: (_ for _ in ()).throw(NotImplementedError("This light tae doesn't support encoding currently"))
-                    self.process_output = lambda image: image
-                    self.memory_used_decode = lambda shape, dtype: (1800 * (max(1, (shape[-3] ** 0.7 * 0.1)) * shape[-2] * shape[-1] * 16 * 16) * model_management.dtype_size(dtype))
-                elif self.latent_channels == 32 and sd["decoder.22.bias"].shape[0] == 12: # lighttae_hv15
-                    self.first_stage_model = comfy.taesd.taehv.TAEHV(latent_channels=self.latent_channels, latent_format=comfy.latent_formats.HunyuanVideo15)
-                    self.process_input = lambda image: (_ for _ in ()).throw(NotImplementedError("This light tae doesn't support encoding currently"))
-                    self.memory_used_decode = lambda shape, dtype: (1200 * (max(1, (shape[-3] ** 0.7 * 0.05)) * shape[-2] * shape[-1] * 32 * 32) * model_management.dtype_size(dtype))
-                else:
-                    if sd["decoder.1.weight"].dtype == torch.float16: # taehv currently only available in float16, so assume it's not lighttaew2_1 as otherwise state dicts are identical
-                        latent_format=comfy.latent_formats.HunyuanVideo
-                    else:
-                        latent_format=None # lighttaew2_1 doesn't need scaling
-                    self.first_stage_model = comfy.taesd.taehv.TAEHV(latent_channels=self.latent_channels, latent_format=latent_format)
-                    self.process_input = self.process_output = lambda image: image
-                    self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
-                    self.upscale_index_formula = (4, 8, 8)
-                    self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)
-                    self.downscale_index_formula = (4, 8, 8)
-                    self.memory_used_encode = lambda shape, dtype: (700 * (max(1, (shape[-3] ** 0.66 * 0.11)) * shape[-2] * shape[-1]) * model_management.dtype_size(dtype))
-                    self.memory_used_decode = lambda shape, dtype: (50 * (max(1, (shape[-3] ** 0.65 * 0.26)) * shape[-2] * shape[-1] * 32 * 32) * model_management.dtype_size(dtype))
             else:
                 logging.warning("WARNING: No VAE weights detected, VAE not initalized.")
                 self.first_stage_model = None
@@ -697,28 +615,17 @@ class VAE:
             raise RuntimeError("ERROR: VAE is invalid: None\n\nIf the VAE is from a checkpoint loader node your checkpoint does not contain a valid VAE.")
 
     def vae_encode_crop_pixels(self, pixels):
-        if self.crop_input:
-            downscale_ratio = self.spacial_compression_encode()
+        if not self.crop_input:
+            return pixels
 
-            dims = pixels.shape[1:-1]
-            for d in range(len(dims)):
-                x = (dims[d] // downscale_ratio) * downscale_ratio
-                x_offset = (dims[d] % downscale_ratio) // 2
-                if x != dims[d]:
-                    pixels = pixels.narrow(d + 1, x_offset, x)
+        downscale_ratio = self.spacial_compression_encode()
 
-        if pixels.shape[-1] > self.output_channels:
-            pixels = pixels[..., :self.output_channels]
-        elif pixels.shape[-1] < self.output_channels:
-            if self.pad_channel_value is not None:
-                if isinstance(self.pad_channel_value, str):
-                    mode = self.pad_channel_value
-                    value = None
-                else:
-                    mode = "constant"
-                    value = self.pad_channel_value
-
-                pixels = torch.nn.functional.pad(pixels, (0, self.output_channels - pixels.shape[-1]), mode=mode, value=value)
+        dims = pixels.shape[1:-1]
+        for d in range(len(dims)):
+            x = (dims[d] // downscale_ratio) * downscale_ratio
+            x_offset = (dims[d] % downscale_ratio) // 2
+            if x != dims[d]:
+                pixels = pixels.narrow(d + 1, x_offset, x)
         return pixels
 
     def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16):
@@ -789,8 +696,6 @@ class VAE:
         self.throw_exception_if_invalid()
         pixel_samples = None
         do_tile = False
-        if self.latent_dim == 2 and samples_in.ndim == 5:
-            samples_in = samples_in[:, :, 0]
         try:
             memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
             model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
@@ -1006,20 +911,12 @@ class CLIPType(Enum):
     OMNIGEN2 = 17
     QWEN_IMAGE = 18
     HUNYUAN_IMAGE = 19
-    HUNYUAN_VIDEO_15 = 20
-    OVIS = 21
-    KANDINSKY5 = 22
-    KANDINSKY5_IMAGE = 23
-    NEWBIE = 24
 
 
 def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}):
     clip_data = []
     for p in ckpt_paths:
-        sd, metadata = comfy.utils.load_torch_file(p, safe_load=True, return_metadata=True)
-        if model_options.get("custom_operations", None) is None:
-            sd, metadata = comfy.utils.convert_old_quants(sd, model_prefix="", metadata=metadata)
-        clip_data.append(sd)
+        clip_data.append(comfy.utils.load_torch_file(p, safe_load=True))
     return load_text_encoder_state_dicts(clip_data, embedding_directory=embedding_directory, clip_type=clip_type, model_options=model_options)
 
 
@@ -1037,12 +934,6 @@ class TEModel(Enum):
     QWEN25_7B = 11
     BYT5_SMALL_GLYPH = 12
     GEMMA_3_4B = 13
-    MISTRAL3_24B = 14
-    MISTRAL3_24B_PRUNED_FLUX2 = 15
-    QWEN3_4B = 16
-    QWEN3_2B = 17
-    JINA_CLIP_2 = 18
-
 
 def detect_te_model(sd):
     if "text_model.encoder.layers.30.mlp.fc1.weight" in sd:
@@ -1051,8 +942,6 @@ def detect_te_model(sd):
         return TEModel.CLIP_H
     if "text_model.encoder.layers.0.mlp.fc1.weight" in sd:
         return TEModel.CLIP_L
-    if "model.encoder.layers.0.mixer.Wqkv.weight" in sd:
-        return TEModel.JINA_CLIP_2
     if "encoder.block.23.layer.1.DenseReluDense.wi_1.weight" in sd:
         weight = sd["encoder.block.23.layer.1.DenseReluDense.wi_1.weight"]
         if weight.shape[-1] == 4096:
@@ -1077,18 +966,6 @@ def detect_te_model(sd):
         if weight.shape[0] == 512:
             return TEModel.QWEN25_7B
     if "model.layers.0.post_attention_layernorm.weight" in sd:
-        weight = sd['model.layers.0.post_attention_layernorm.weight']
-        if 'model.layers.0.self_attn.q_norm.weight' in sd:
-            if weight.shape[0] == 2560:
-                return TEModel.QWEN3_4B
-            elif weight.shape[0] == 2048:
-                return TEModel.QWEN3_2B
-        if weight.shape[0] == 5120:
-            if "model.layers.39.post_attention_layernorm.weight" in sd:
-                return TEModel.MISTRAL3_24B
-            else:
-                return TEModel.MISTRAL3_24B_PRUNED_FLUX2
-
         return TEModel.LLAMA3_8
     return None
 
@@ -1138,7 +1015,7 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
                 clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=False, clip_g=True, t5=False)
                 clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
             elif clip_type == CLIPType.HIDREAM:
-                clip_target.clip = comfy.text_encoders.hidream.hidream_clip(clip_l=False, clip_g=True, t5=False, llama=False, dtype_t5=None, dtype_llama=None)
+                clip_target.clip = comfy.text_encoders.hidream.hidream_clip(clip_l=False, clip_g=True, t5=False, llama=False, dtype_t5=None, dtype_llama=None, t5xxl_scaled_fp8=None, llama_scaled_fp8=None)
                 clip_target.tokenizer = comfy.text_encoders.hidream.HiDreamTokenizer
             else:
                 clip_target.clip = sdxl_clip.SDXLRefinerClipModel
@@ -1162,7 +1039,7 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
                 tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
             elif clip_type == CLIPType.HIDREAM:
                 clip_target.clip = comfy.text_encoders.hidream.hidream_clip(**t5xxl_detect(clip_data),
-                                                                        clip_l=False, clip_g=False, t5=True, llama=False, dtype_llama=None)
+                                                                        clip_l=False, clip_g=False, t5=True, llama=False, dtype_llama=None, llama_scaled_fp8=None)
                 clip_target.tokenizer = comfy.text_encoders.hidream.HiDreamTokenizer
             else: #CLIPType.MOCHI
                 clip_target.clip = comfy.text_encoders.genmo.mochi_te(**t5xxl_detect(clip_data))
@@ -1191,7 +1068,7 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
         elif te_model == TEModel.LLAMA3_8:
             clip_target.clip = comfy.text_encoders.hidream.hidream_clip(**llama_detect(clip_data),
-                                                                        clip_l=False, clip_g=False, t5=False, llama=True, dtype_t5=None)
+                                                                        clip_l=False, clip_g=False, t5=False, llama=True, dtype_t5=None, t5xxl_scaled_fp8=None)
             clip_target.tokenizer = comfy.text_encoders.hidream.HiDreamTokenizer
         elif te_model == TEModel.QWEN25_3B:
             clip_target.clip = comfy.text_encoders.omnigen2.te(**llama_detect(clip_data))
@@ -1203,26 +1080,13 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             else:
                 clip_target.clip = comfy.text_encoders.qwen_image.te(**llama_detect(clip_data))
                 clip_target.tokenizer = comfy.text_encoders.qwen_image.QwenImageTokenizer
-        elif te_model == TEModel.MISTRAL3_24B or te_model == TEModel.MISTRAL3_24B_PRUNED_FLUX2:
-            clip_target.clip = comfy.text_encoders.flux.flux2_te(**llama_detect(clip_data), pruned=te_model == TEModel.MISTRAL3_24B_PRUNED_FLUX2)
-            clip_target.tokenizer = comfy.text_encoders.flux.Flux2Tokenizer
-            tokenizer_data["tekken_model"] = clip_data[0].get("tekken_model", None)
-        elif te_model == TEModel.QWEN3_4B:
-            clip_target.clip = comfy.text_encoders.z_image.te(**llama_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.z_image.ZImageTokenizer
-        elif te_model == TEModel.QWEN3_2B:
-            clip_target.clip = comfy.text_encoders.ovis.te(**llama_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.ovis.OvisTokenizer
-        elif te_model == TEModel.JINA_CLIP_2:
-            clip_target.clip = comfy.text_encoders.jina_clip_2.JinaClip2TextModelWrapper
-            clip_target.tokenizer = comfy.text_encoders.jina_clip_2.JinaClip2TokenizerWrapper
         else:
             # clip_l
             if clip_type == CLIPType.SD3:
                 clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=True, clip_g=False, t5=False)
                 clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
             elif clip_type == CLIPType.HIDREAM:
-                clip_target.clip = comfy.text_encoders.hidream.hidream_clip(clip_l=True, clip_g=False, t5=False, llama=False, dtype_t5=None, dtype_llama=None)
+                clip_target.clip = comfy.text_encoders.hidream.hidream_clip(clip_l=True, clip_g=False, t5=False, llama=False, dtype_t5=None, dtype_llama=None, t5xxl_scaled_fp8=None, llama_scaled_fp8=None)
                 clip_target.tokenizer = comfy.text_encoders.hidream.HiDreamTokenizer
             else:
                 clip_target.clip = sd1_clip.SD1ClipModel
@@ -1262,26 +1126,6 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
         elif clip_type == CLIPType.HUNYUAN_IMAGE:
             clip_target.clip = comfy.text_encoders.hunyuan_image.te(**llama_detect(clip_data))
             clip_target.tokenizer = comfy.text_encoders.hunyuan_image.HunyuanImageTokenizer
-        elif clip_type == CLIPType.HUNYUAN_VIDEO_15:
-            clip_target.clip = comfy.text_encoders.hunyuan_image.te(**llama_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.hunyuan_video.HunyuanVideo15Tokenizer
-        elif clip_type == CLIPType.KANDINSKY5:
-            clip_target.clip = comfy.text_encoders.kandinsky5.te(**llama_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.kandinsky5.Kandinsky5Tokenizer
-        elif clip_type == CLIPType.KANDINSKY5_IMAGE:
-            clip_target.clip = comfy.text_encoders.kandinsky5.te(**llama_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.kandinsky5.Kandinsky5TokenizerImage
-        elif clip_type == CLIPType.NEWBIE:
-            clip_target.clip = comfy.text_encoders.newbie.te(**llama_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.newbie.NewBieTokenizer
-            if "model.layers.0.self_attn.q_norm.weight" in clip_data[0]:
-                clip_data_gemma = clip_data[0]
-                clip_data_jina = clip_data[1]
-            else:
-                clip_data_gemma = clip_data[1]
-                clip_data_jina = clip_data[0]
-            tokenizer_data["gemma_spiece_model"] = clip_data_gemma.get("spiece_model", None)
-            tokenizer_data["jina_spiece_model"] = clip_data_jina.get("spiece_model", None)
         else:
             clip_target.clip = sdxl_clip.SDXLClipModel
             clip_target.tokenizer = sdxl_clip.SDXLTokenizer
@@ -1297,7 +1141,14 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
         parameters += comfy.utils.calculate_parameters(c)
         tokenizer_data, model_options = comfy.text_encoders.long_clipl.model_options_long_clip(c, tokenizer_data, model_options)
 
-    clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters, tokenizer_data=tokenizer_data, state_dict=clip_data, model_options=model_options)
+    clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters, tokenizer_data=tokenizer_data, model_options=model_options)
+    for c in clip_data:
+        m, u = clip.load_sd(c)
+        if len(m) > 0:
+            logging.warning("clip missing: {}".format(m))
+
+        if len(u) > 0:
+            logging.debug("clip unexpected: {}".format(u))
     return clip
 
 def load_gligen(ckpt_path):
@@ -1356,10 +1207,6 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
     weight_dtype = comfy.utils.weight_dtype(sd, diffusion_model_prefix)
     load_device = model_management.get_torch_device()
 
-    custom_operations = model_options.get("custom_operations", None)
-    if custom_operations is None:
-        sd, metadata = comfy.utils.convert_old_quants(sd, diffusion_model_prefix, metadata=metadata)
-
     model_config = model_detection.model_config_from_unet(sd, diffusion_model_prefix, metadata=metadata)
     if model_config is None:
         logging.warning("Warning, This is not a checkpoint file, trying to load it as a diffusion model only.")
@@ -1368,22 +1215,18 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
             return None
         return (diffusion_model, None, VAE(sd={}), None)  # The VAE object is there to throw an exception if it's actually used'
 
+
     unet_weight_dtype = list(model_config.supported_inference_dtypes)
-    if model_config.quant_config is not None:
+    if model_config.scaled_fp8 is not None:
         weight_dtype = None
 
-    if custom_operations is not None:
-        model_config.custom_operations = custom_operations
-
+    model_config.custom_operations = model_options.get("custom_operations", None)
     unet_dtype = model_options.get("dtype", model_options.get("weight_dtype", None))
 
     if unet_dtype is None:
         unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype, weight_dtype=weight_dtype)
 
-    if model_config.quant_config is not None:
-        manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
-    else:
-        manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
+    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.clip_vision_prefix is not None:
@@ -1401,33 +1244,22 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
         vae = VAE(sd=vae_sd, metadata=metadata)
 
     if output_clip:
-        if te_model_options.get("custom_operations", None) is None:
-            scaled_fp8_list = []
-            for k in list(sd.keys()):  # Convert scaled fp8 to mixed ops
-                if k.endswith(".scaled_fp8"):
-                    scaled_fp8_list.append(k[:-len("scaled_fp8")])
-
-            if len(scaled_fp8_list) > 0:
-                out_sd = {}
-                for k in sd:
-                    skip = False
-                    for pref in scaled_fp8_list:
-                        skip = skip or k.startswith(pref)
-                    if not skip:
-                        out_sd[k] = sd[k]
-
-                for pref in scaled_fp8_list:
-                    quant_sd, qmetadata = comfy.utils.convert_old_quants(sd, pref, metadata={})
-                    for k in quant_sd:
-                        out_sd[k] = quant_sd[k]
-                    sd = out_sd
-
         clip_target = model_config.clip_target(state_dict=sd)
         if clip_target is not None:
             clip_sd = model_config.process_clip_state_dict(sd)
             if len(clip_sd) > 0:
                 parameters = comfy.utils.calculate_parameters(clip_sd)
-                clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd, parameters=parameters, state_dict=clip_sd, model_options=te_model_options)
+                clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd, parameters=parameters, model_options=te_model_options)
+                m, u = clip.load_sd(clip_sd, full_model=True)
+                if len(m) > 0:
+                    m_filter = list(filter(lambda a: ".logit_scale" not in a and ".transformer.text_projection.weight" not in a, m))
+                    if len(m_filter) > 0:
+                        logging.warning("clip missing: {}".format(m))
+                    else:
+                        logging.debug("clip missing: {}".format(m))
+
+                if len(u) > 0:
+                    logging.debug("clip unexpected {}:".format(u))
             else:
                 logging.warning("no CLIP/text encoder weights in checkpoint, the text encoder model will not be loaded.")
 
@@ -1474,9 +1306,6 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None):
     if len(temp_sd) > 0:
         sd = temp_sd
 
-    custom_operations = model_options.get("custom_operations", None)
-    if custom_operations is None:
-        sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata)
     parameters = comfy.utils.calculate_parameters(sd)
     weight_dtype = comfy.utils.weight_dtype(sd)
 
@@ -1507,7 +1336,7 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None):
 
     offload_device = model_management.unet_offload_device()
     unet_weight_dtype = list(model_config.supported_inference_dtypes)
-    if model_config.quant_config is not None:
+    if model_config.scaled_fp8 is not None:
         weight_dtype = None
 
     if dtype is None:
@@ -1515,15 +1344,12 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None):
     else:
         unet_dtype = dtype
 
-    if model_config.quant_config is not None:
+    if model_config.layer_quant_config is not None:
         manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
     else:
         manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
     model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
-
-    if custom_operations is not None:
-        model_config.custom_operations = custom_operations
-
+    model_config.custom_operations = model_options.get("custom_operations", model_config.custom_operations)
     if model_options.get("fp8_optimizations", False):
         model_config.optimizations["fp8"] = True
 
@@ -1562,9 +1388,6 @@ def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, m
     if vae is not None:
         vae_sd = vae.get_sd()
 
-    if metadata is None:
-        metadata = {}
-
     model_management.load_models_gpu(load_models, force_patch_weights=True)
     clip_vision_sd = clip_vision.get_sd() if clip_vision is not None else None
     sd = model.model.state_dict_for_saving(clip_sd, vae_sd, clip_vision_sd)

comfy/sd1_clip.py

@@ -90,6 +90,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
                  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, return_attention_masks=False, model_options={}):  # clip-vit-base-patch32
         super().__init__()
+        assert layer in self.LAYERS
 
         if textmodel_json_config is None:
             textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_clip_config.json")
@@ -107,17 +108,19 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
             config[k] = v
 
         operations = model_options.get("custom_operations", None)
-        quant_config = model_options.get("quantization_metadata", None)
+        scaled_fp8 = None
 
         if operations is None:
-            if quant_config is not None:
-                operations = comfy.ops.mixed_precision_ops(quant_config, dtype, full_precision_mm=True)
-                logging.info("Using MixedPrecisionOps for text encoder")
+            scaled_fp8 = model_options.get("scaled_fp8", None)
+            if scaled_fp8 is not None:
+                operations = comfy.ops.scaled_fp8_ops(fp8_matrix_mult=False, override_dtype=scaled_fp8)
             else:
                 operations = comfy.ops.manual_cast
 
         self.operations = operations
         self.transformer = model_class(config, dtype, device, self.operations)
+        if scaled_fp8 is not None:
+            self.transformer.scaled_fp8 = torch.nn.Parameter(torch.tensor([], dtype=scaled_fp8))
 
         self.num_layers = self.transformer.num_layers
 
@@ -135,7 +138,6 @@ 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
-        self.execution_device = None
 
         if layer == "hidden":
             assert layer_idx is not None
@@ -152,8 +154,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
     def set_clip_options(self, options):
         layer_idx = options.get("layer", self.layer_idx)
         self.return_projected_pooled = options.get("projected_pooled", self.return_projected_pooled)
-        self.execution_device = options.get("execution_device", self.execution_device)
-        if isinstance(self.layer, list) or self.layer == "all":
+        if self.layer == "all":
             pass
         elif layer_idx is None or abs(layer_idx) > self.num_layers:
             self.layer = "last"
@@ -165,7 +166,6 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
         self.layer = self.options_default[0]
         self.layer_idx = self.options_default[1]
         self.return_projected_pooled = self.options_default[2]
-        self.execution_device = None
 
     def process_tokens(self, tokens, device):
         end_token = self.special_tokens.get("end", None)
@@ -249,20 +249,14 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
         return torch.cat(embeds_out), torch.tensor(attention_masks, device=device, dtype=torch.long), num_tokens, embeds_info
 
     def forward(self, tokens):
-        if self.execution_device is None:
-            device = self.transformer.get_input_embeddings().weight.device
-        else:
-            device = self.execution_device
-
+        device = self.transformer.get_input_embeddings().weight.device
         embeds, attention_mask, num_tokens, embeds_info = self.process_tokens(tokens, device)
 
         attention_mask_model = None
         if self.enable_attention_masks:
             attention_mask_model = attention_mask
 
-        if isinstance(self.layer, list):
-            intermediate_output = self.layer
-        elif self.layer == "all":
+        if self.layer == "all":
             intermediate_output = "all"
         else:
             intermediate_output = self.layer_idx
@@ -466,7 +460,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, has_end_token=True, pad_to_max_length=True, min_length=None, pad_token=None, end_token=None, min_padding=None, pad_left=False, disable_weights=False, tokenizer_data={}, tokenizer_args={}):
+    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, has_end_token=True, pad_to_max_length=True, min_length=None, pad_token=None, end_token=None, min_padding=None, tokenizer_data={}, tokenizer_args={}):
         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, **tokenizer_args)
@@ -474,7 +468,6 @@ class SDTokenizer:
         self.min_length = tokenizer_data.get("{}_min_length".format(embedding_key), min_length)
         self.end_token = None
         self.min_padding = min_padding
-        self.pad_left = pad_left
 
         empty = self.tokenizer('')["input_ids"]
         self.tokenizer_adds_end_token = has_end_token
@@ -513,8 +506,6 @@ class SDTokenizer:
         self.embedding_size = embedding_size
         self.embedding_key = embedding_key
 
-        self.disable_weights = disable_weights
-
     def _try_get_embedding(self, embedding_name:str):
         '''
         Takes a potential embedding name and tries to retrieve it.
@@ -531,12 +522,6 @@ class SDTokenizer:
                 return (embed, "{} {}".format(embedding_name[len(stripped):], leftover))
         return (embed, leftover)
 
-    def pad_tokens(self, tokens, amount):
-        if self.pad_left:
-            for i in range(amount):
-                tokens.insert(0, (self.pad_token, 1.0, 0))
-        else:
-            tokens.extend([(self.pad_token, 1.0, 0)] * amount)
 
     def tokenize_with_weights(self, text:str, return_word_ids=False, tokenizer_options={}, **kwargs):
         '''
@@ -549,7 +534,7 @@ class SDTokenizer:
         min_padding = tokenizer_options.get("{}_min_padding".format(self.embedding_key), self.min_padding)
 
         text = escape_important(text)
-        if kwargs.get("disable_weights", self.disable_weights):
+        if kwargs.get("disable_weights", False):
             parsed_weights = [(text, 1.0)]
         else:
             parsed_weights = token_weights(text, 1.0)
@@ -615,7 +600,7 @@ class SDTokenizer:
                         if self.end_token is not None:
                             batch.append((self.end_token, 1.0, 0))
                         if self.pad_to_max_length:
-                            self.pad_tokens(batch, remaining_length)
+                            batch.extend([(self.pad_token, 1.0, 0)] * (remaining_length))
                     #start new batch
                     batch = []
                     if self.start_token is not None:
@@ -629,11 +614,11 @@ class SDTokenizer:
         if self.end_token is not None:
             batch.append((self.end_token, 1.0, 0))
         if min_padding is not None:
-            self.pad_tokens(batch, min_padding)
+            batch.extend([(self.pad_token, 1.0, 0)] * min_padding)
         if self.pad_to_max_length and len(batch) < self.max_length:
-            self.pad_tokens(batch, self.max_length - len(batch))
+            batch.extend([(self.pad_token, 1.0, 0)] * (self.max_length - len(batch)))
         if min_length is not None and len(batch) < min_length:
-            self.pad_tokens(batch, min_length - len(batch))
+            batch.extend([(self.pad_token, 1.0, 0)] * (min_length - len(batch)))
 
         if not return_word_ids:
             batched_tokens = [[(t, w) for t, w,_ in x] for x in batched_tokens]

comfy/supported_models.py

@@ -21,14 +21,11 @@ import comfy.text_encoders.ace
 import comfy.text_encoders.omnigen2
 import comfy.text_encoders.qwen_image
 import comfy.text_encoders.hunyuan_image
-import comfy.text_encoders.kandinsky5
-import comfy.text_encoders.z_image
 
 from . import supported_models_base
 from . import latent_formats
 
 from . import diffusers_convert
-import comfy.model_management
 
 class SD15(supported_models_base.BASE):
     unet_config = {
@@ -542,7 +539,7 @@ class SD3(supported_models_base.BASE):
     unet_extra_config = {}
     latent_format = latent_formats.SD3
 
-    memory_usage_factor = 1.6
+    memory_usage_factor = 1.2
 
     text_encoder_key_prefix = ["text_encoders."]
 
@@ -744,37 +741,6 @@ class FluxSchnell(Flux):
         out = model_base.Flux(self, model_type=model_base.ModelType.FLOW, device=device)
         return out
 
-class Flux2(Flux):
-    unet_config = {
-        "image_model": "flux2",
-    }
-
-    sampling_settings = {
-        "shift": 2.02,
-    }
-
-    unet_extra_config = {}
-    latent_format = latent_formats.Flux2
-
-    supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
-
-    vae_key_prefix = ["vae."]
-    text_encoder_key_prefix = ["text_encoders."]
-
-    def __init__(self, unet_config):
-        super().__init__(unet_config)
-        self.memory_usage_factor = self.memory_usage_factor * (2.0 * 2.0) * 2.36
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.Flux2(self, device=device)
-        return out
-
-    def clip_target(self, state_dict={}):
-        return None # TODO
-        pref = self.text_encoder_key_prefix[0]
-        t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))
-        return supported_models_base.ClipTarget(comfy.text_encoders.flux.FluxTokenizer, comfy.text_encoders.flux.flux_clip(**t5_detect))
-
 class GenmoMochi(supported_models_base.BASE):
     unet_config = {
         "image_model": "mochi_preview",
@@ -966,7 +932,7 @@ class CosmosT2IPredict2(supported_models_base.BASE):
 
     def __init__(self, unet_config):
         super().__init__(unet_config)
-        self.memory_usage_factor = (unet_config.get("model_channels", 2048) / 2048) * 0.95
+        self.memory_usage_factor = (unet_config.get("model_channels", 2048) / 2048) * 0.9
 
     def get_model(self, state_dict, prefix="", device=None):
         out = model_base.CosmosPredict2(self, device=device)
@@ -997,7 +963,7 @@ class Lumina2(supported_models_base.BASE):
         "shift": 6.0,
     }
 
-    memory_usage_factor = 1.4
+    memory_usage_factor = 1.2
 
     unet_extra_config = {}
     latent_format = latent_formats.Flux
@@ -1016,32 +982,6 @@ class Lumina2(supported_models_base.BASE):
         hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}gemma2_2b.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.lumina2.LuminaTokenizer, comfy.text_encoders.lumina2.te(**hunyuan_detect))
 
-class ZImage(Lumina2):
-    unet_config = {
-        "image_model": "lumina2",
-        "dim": 3840,
-    }
-
-    sampling_settings = {
-        "multiplier": 1.0,
-        "shift": 3.0,
-    }
-
-    memory_usage_factor = 2.0
-
-    supported_inference_dtypes = [torch.bfloat16, torch.float32]
-
-    def __init__(self, unet_config):
-        super().__init__(unet_config)
-        if comfy.model_management.extended_fp16_support():
-            self.supported_inference_dtypes = self.supported_inference_dtypes.copy()
-            self.supported_inference_dtypes.insert(1, torch.float16)
-
-    def clip_target(self, state_dict={}):
-        pref = self.text_encoder_key_prefix[0]
-        hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3_4b.transformer.".format(pref))
-        return supported_models_base.ClipTarget(comfy.text_encoders.z_image.ZImageTokenizer, comfy.text_encoders.z_image.te(**hunyuan_detect))
-
 class WAN21_T2V(supported_models_base.BASE):
     unet_config = {
         "image_model": "wan2.1",
@@ -1296,7 +1236,7 @@ class ChromaRadiance(Chroma):
     latent_format = comfy.latent_formats.ChromaRadiance
 
     # Pixel-space model, no spatial compression for model input.
-    memory_usage_factor = 0.044
+    memory_usage_factor = 0.038
 
     def get_model(self, state_dict, prefix="", device=None):
         return model_base.ChromaRadiance(self, device=device)
@@ -1339,7 +1279,7 @@ class Omnigen2(supported_models_base.BASE):
         "shift": 2.6,
     }
 
-    memory_usage_factor = 1.95 #TODO
+    memory_usage_factor = 1.65 #TODO
 
     unet_extra_config = {}
     latent_format = latent_formats.Flux
@@ -1404,7 +1344,7 @@ class HunyuanImage21(HunyuanVideo):
 
     latent_format = latent_formats.HunyuanImage21
 
-    memory_usage_factor = 8.7
+    memory_usage_factor = 7.7
 
     supported_inference_dtypes = [torch.bfloat16, torch.float32]
 
@@ -1434,108 +1374,6 @@ class HunyuanImage21Refiner(HunyuanVideo):
         out = model_base.HunyuanImage21Refiner(self, device=device)
         return out
 
-class HunyuanVideo15(HunyuanVideo):
-    unet_config = {
-        "image_model": "hunyuan_video",
-        "vision_in_dim": 1152,
-    }
-
-    sampling_settings = {
-        "shift": 7.0,
-    }
-    memory_usage_factor = 4.0 #TODO
-    supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
-
-    latent_format = latent_formats.HunyuanVideo15
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.HunyuanVideo15(self, device=device)
-        return out
-
-    def clip_target(self, state_dict={}):
-        pref = self.text_encoder_key_prefix[0]
-        hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
-        return supported_models_base.ClipTarget(comfy.text_encoders.hunyuan_video.HunyuanVideo15Tokenizer, comfy.text_encoders.hunyuan_image.te(**hunyuan_detect))
-
-
-class HunyuanVideo15_SR_Distilled(HunyuanVideo):
-    unet_config = {
-        "image_model": "hunyuan_video",
-        "vision_in_dim": 1152,
-        "in_channels": 98,
-    }
-
-    sampling_settings = {
-        "shift": 2.0,
-    }
-    memory_usage_factor = 4.0 #TODO
-    supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
-
-    latent_format = latent_formats.HunyuanVideo15
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.HunyuanVideo15_SR_Distilled(self, device=device)
-        return out
-
-    def clip_target(self, state_dict={}):
-        pref = self.text_encoder_key_prefix[0]
-        hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
-        return supported_models_base.ClipTarget(comfy.text_encoders.hunyuan_video.HunyuanVideo15Tokenizer, comfy.text_encoders.hunyuan_image.te(**hunyuan_detect))
-
-
-class Kandinsky5(supported_models_base.BASE):
-    unet_config = {
-        "image_model": "kandinsky5",
-    }
-
-    sampling_settings = {
-        "shift": 10.0,
-    }
-
-    unet_extra_config = {}
-    latent_format = latent_formats.HunyuanVideo
-
-    memory_usage_factor = 1.25 #TODO
-
-    supported_inference_dtypes = [torch.bfloat16, torch.float32]
-
-    vae_key_prefix = ["vae."]
-    text_encoder_key_prefix = ["text_encoders."]
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.Kandinsky5(self, device=device)
-        return out
-
-    def clip_target(self, state_dict={}):
-        pref = self.text_encoder_key_prefix[0]
-        hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
-        return supported_models_base.ClipTarget(comfy.text_encoders.kandinsky5.Kandinsky5Tokenizer, comfy.text_encoders.kandinsky5.te(**hunyuan_detect))
-
-
-class Kandinsky5Image(Kandinsky5):
-    unet_config = {
-        "image_model": "kandinsky5",
-        "model_dim": 2560,
-        "visual_embed_dim": 64,
-    }
-
-    sampling_settings = {
-        "shift": 3.0,
-    }
-
-    latent_format = latent_formats.Flux
-    memory_usage_factor = 1.25 #TODO
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.Kandinsky5Image(self, device=device)
-        return out
-
-    def clip_target(self, state_dict={}):
-        pref = self.text_encoder_key_prefix[0]
-        hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
-        return supported_models_base.ClipTarget(comfy.text_encoders.kandinsky5.Kandinsky5TokenizerImage, comfy.text_encoders.kandinsky5.te(**hunyuan_detect))
-
-
-models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanVideo15_SR_Distilled, HunyuanVideo15, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, ZImage, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, Omnigen2, QwenImage, Flux2, Kandinsky5Image, Kandinsky5]
+models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, Omnigen2, QwenImage]
 
 models += [SVD_img2vid]

comfy/supported_models_base.py

@@ -17,7 +17,6 @@
 """
 
 import torch
-import logging
 from . import model_base
 from . import utils
 from . import latent_formats
@@ -50,7 +49,8 @@ class BASE:
 
     manual_cast_dtype = None
     custom_operations = None
-    quant_config = None  # quantization configuration for mixed precision
+    scaled_fp8 = None
+    layer_quant_config = None  # Per-layer quantization configuration for mixed precision
     optimizations = {"fp8": False}
 
     @classmethod
@@ -118,7 +118,3 @@ class BASE:
     def set_inference_dtype(self, dtype, manual_cast_dtype):
         self.unet_config['dtype'] = dtype
         self.manual_cast_dtype = manual_cast_dtype
-
-    def __getattr__(self, name):
-        logging.warning("\nWARNING, you accessed {} from the model config object which doesn't exist. Please fix your code.\n".format(name))
-        return None

comfy/taesd/taehv.py

@@ -1,171 +0,0 @@
-# Tiny AutoEncoder for HunyuanVideo and WanVideo https://github.com/madebyollin/taehv
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from tqdm.auto import tqdm
-from collections import namedtuple, deque
-
-import comfy.ops
-operations=comfy.ops.disable_weight_init
-
-DecoderResult = namedtuple("DecoderResult", ("frame", "memory"))
-TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index"))
-
-def conv(n_in, n_out, **kwargs):
-    return operations.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
-
-class Clamp(nn.Module):
-    def forward(self, x):
-        return torch.tanh(x / 3) * 3
-
-class MemBlock(nn.Module):
-    def __init__(self, n_in, n_out, act_func):
-        super().__init__()
-        self.conv = nn.Sequential(conv(n_in * 2, n_out), act_func, conv(n_out, n_out), act_func, conv(n_out, n_out))
-        self.skip = operations.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
-        self.act = act_func
-    def forward(self, x, past):
-        return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x))
-
-class TPool(nn.Module):
-    def __init__(self, n_f, stride):
-        super().__init__()
-        self.stride = stride
-        self.conv = operations.Conv2d(n_f*stride,n_f, 1, bias=False)
-    def forward(self, x):
-        _NT, C, H, W = x.shape
-        return self.conv(x.reshape(-1, self.stride * C, H, W))
-
-class TGrow(nn.Module):
-    def __init__(self, n_f, stride):
-        super().__init__()
-        self.stride = stride
-        self.conv = operations.Conv2d(n_f, n_f*stride, 1, bias=False)
-    def forward(self, x):
-        _NT, C, H, W = x.shape
-        x = self.conv(x)
-        return x.reshape(-1, C, H, W)
-
-def apply_model_with_memblocks(model, x, parallel, show_progress_bar):
-
-    B, T, C, H, W = x.shape
-    if parallel:
-        x = x.reshape(B*T, C, H, W)
-        # parallel over input timesteps, iterate over blocks
-        for b in tqdm(model, disable=not show_progress_bar):
-            if isinstance(b, MemBlock):
-                BT, C, H, W = x.shape
-                T = BT // B
-                _x = x.reshape(B, T, C, H, W)
-                mem = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape)
-                x = b(x, mem)
-            else:
-                x = b(x)
-        BT, C, H, W = x.shape
-        T = BT // B
-        x = x.view(B, T, C, H, W)
-    else:
-        out = []
-        work_queue = deque([TWorkItem(xt, 0) for t, xt in enumerate(x.reshape(B, T * C, H, W).chunk(T, dim=1))])
-        progress_bar = tqdm(range(T), disable=not show_progress_bar)
-        mem = [None] * len(model)
-        while work_queue:
-            xt, i = work_queue.popleft()
-            if i == 0:
-                progress_bar.update(1)
-            if i == len(model):
-                out.append(xt)
-                del xt
-            else:
-                b = model[i]
-                if isinstance(b, MemBlock):
-                    if mem[i] is None:
-                        xt_new = b(xt, xt * 0)
-                        mem[i] = xt.detach().clone()
-                    else:
-                        xt_new = b(xt, mem[i])
-                        mem[i] = xt.detach().clone()
-                    del xt
-                    work_queue.appendleft(TWorkItem(xt_new, i+1))
-                elif isinstance(b, TPool):
-                    if mem[i] is None:
-                        mem[i] = []
-                    mem[i].append(xt.detach().clone())
-                    if len(mem[i]) == b.stride:
-                        B, C, H, W = xt.shape
-                        xt = b(torch.cat(mem[i], 1).view(B*b.stride, C, H, W))
-                        mem[i] = []
-                        work_queue.appendleft(TWorkItem(xt, i+1))
-                elif isinstance(b, TGrow):
-                    xt = b(xt)
-                    NT, C, H, W = xt.shape
-                    for xt_next in reversed(xt.view(B, b.stride*C, H, W).chunk(b.stride, 1)):
-                        work_queue.appendleft(TWorkItem(xt_next, i+1))
-                    del xt
-                else:
-                    xt = b(xt)
-                    work_queue.appendleft(TWorkItem(xt, i+1))
-        progress_bar.close()
-        x = torch.stack(out, 1)
-    return x
-
-
-class TAEHV(nn.Module):
-    def __init__(self, latent_channels, parallel=False, decoder_time_upscale=(True, True), decoder_space_upscale=(True, True, True), latent_format=None, show_progress_bar=True):
-        super().__init__()
-        self.image_channels = 3
-        self.patch_size = 1
-        self.latent_channels = latent_channels
-        self.parallel = parallel
-        self.latent_format = latent_format
-        self.show_progress_bar = show_progress_bar
-        self.process_in = latent_format().process_in if latent_format is not None else (lambda x: x)
-        self.process_out = latent_format().process_out if latent_format is not None else (lambda x: x)
-        if self.latent_channels in [48, 32]: # Wan 2.2 and HunyuanVideo1.5
-            self.patch_size = 2
-        if self.latent_channels == 32: # HunyuanVideo1.5
-            act_func = nn.LeakyReLU(0.2, inplace=True)
-        else: # HunyuanVideo, Wan 2.1
-            act_func = nn.ReLU(inplace=True)
-
-        self.encoder = nn.Sequential(
-            conv(self.image_channels*self.patch_size**2, 64), act_func,
-            TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64, act_func), MemBlock(64, 64, act_func), MemBlock(64, 64, act_func),
-            TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64, act_func), MemBlock(64, 64, act_func), MemBlock(64, 64, act_func),
-            TPool(64, 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64, act_func), MemBlock(64, 64, act_func), MemBlock(64, 64, act_func),
-            conv(64, self.latent_channels),
-        )
-        n_f = [256, 128, 64, 64]
-        self.frames_to_trim = 2**sum(decoder_time_upscale) - 1
-        self.decoder = nn.Sequential(
-            Clamp(), conv(self.latent_channels, n_f[0]), act_func,
-            MemBlock(n_f[0], n_f[0], act_func), MemBlock(n_f[0], n_f[0], act_func), MemBlock(n_f[0], n_f[0], act_func), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False),
-            MemBlock(n_f[1], n_f[1], act_func), MemBlock(n_f[1], n_f[1], act_func), MemBlock(n_f[1], n_f[1], act_func), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False),
-            MemBlock(n_f[2], n_f[2], act_func), MemBlock(n_f[2], n_f[2], act_func), MemBlock(n_f[2], n_f[2], act_func), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), conv(n_f[2], n_f[3], bias=False),
-            act_func, conv(n_f[3], self.image_channels*self.patch_size**2),
-        )
-        @property
-        def show_progress_bar(self):
-            return self._show_progress_bar
-
-        @show_progress_bar.setter
-        def show_progress_bar(self, value):
-            self._show_progress_bar = value
-
-    def encode(self, x, **kwargs):
-        if self.patch_size > 1: x = F.pixel_unshuffle(x, self.patch_size)
-        x = x.movedim(2, 1)  # [B, C, T, H, W] -> [B, T, C, H, W]
-        if x.shape[1] % 4 != 0:
-            # pad at end to multiple of 4
-            n_pad = 4 - x.shape[1] % 4
-            padding = x[:, -1:].repeat_interleave(n_pad, dim=1)
-            x = torch.cat([x, padding], 1)
-        x = apply_model_with_memblocks(self.encoder, x, self.parallel, self.show_progress_bar).movedim(2, 1)
-        return self.process_out(x)
-
-    def decode(self, x, **kwargs):
-        x = self.process_in(x).movedim(2, 1)  # [B, C, T, H, W] -> [B, T, C, H, W]
-        x = apply_model_with_memblocks(self.decoder, x, self.parallel, self.show_progress_bar)
-        if self.patch_size > 1: x = F.pixel_shuffle(x, self.patch_size)
-        return x[:, self.frames_to_trim:].movedim(2, 1)

comfy/text_encoders/cosmos.py

@@ -7,10 +7,10 @@ from transformers import T5TokenizerFast
 class T5XXLModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
         textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_old_config_xxl.json")
-        t5xxl_quantization_metadata = model_options.get("t5xxl_quantization_metadata", None)
-        if t5xxl_quantization_metadata is not None:
+        t5xxl_scaled_fp8 = model_options.get("t5xxl_scaled_fp8", None)
+        if t5xxl_scaled_fp8 is not None:
             model_options = model_options.copy()
-            model_options["quantization_metadata"] = t5xxl_quantization_metadata
+            model_options["scaled_fp8"] = t5xxl_scaled_fp8
 
         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=attention_mask, return_attention_masks=attention_mask, zero_out_masked=attention_mask, model_options=model_options)
 
@@ -30,12 +30,12 @@ class CosmosT5Tokenizer(sd1_clip.SD1Tokenizer):
         super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, clip_name="t5xxl", tokenizer=T5XXLTokenizer)
 
 
-def te(dtype_t5=None, t5_quantization_metadata=None):
+def te(dtype_t5=None, t5xxl_scaled_fp8=None):
     class CosmosTEModel_(CosmosT5XXL):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if t5_quantization_metadata is not None:
+            if t5xxl_scaled_fp8 is not None and "t5xxl_scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["t5xxl_quantization_metadata"] = t5_quantization_metadata
+                model_options["t5xxl_scaled_fp8"] = t5xxl_scaled_fp8
             if dtype is None:
                 dtype = dtype_t5
             super().__init__(device=device, dtype=dtype, model_options=model_options)

comfy/text_encoders/flux.py

@@ -1,13 +1,10 @@
 from comfy import sd1_clip
 import comfy.text_encoders.t5
 import comfy.text_encoders.sd3_clip
-import comfy.text_encoders.llama
 import comfy.model_management
-from transformers import T5TokenizerFast, LlamaTokenizerFast
+from transformers import T5TokenizerFast
 import torch
 import os
-import json
-import base64
 
 class T5XXLTokenizer(sd1_clip.SDTokenizer):
     def __init__(self, embedding_directory=None, tokenizer_data={}):
@@ -63,112 +60,11 @@ class FluxClipModel(torch.nn.Module):
         else:
             return self.t5xxl.load_sd(sd)
 
-def flux_clip(dtype_t5=None, t5_quantization_metadata=None):
+def flux_clip(dtype_t5=None, t5xxl_scaled_fp8=None):
     class FluxClipModel_(FluxClipModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if t5_quantization_metadata is not None:
+            if t5xxl_scaled_fp8 is not None and "t5xxl_scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["t5xxl_quantization_metadata"] = t5_quantization_metadata
+                model_options["t5xxl_scaled_fp8"] = t5xxl_scaled_fp8
             super().__init__(dtype_t5=dtype_t5, device=device, dtype=dtype, model_options=model_options)
     return FluxClipModel_
-
-def load_mistral_tokenizer(data):
-    if torch.is_tensor(data):
-        data = data.numpy().tobytes()
-
-    try:
-        from transformers.integrations.mistral import MistralConverter
-    except ModuleNotFoundError:
-        from transformers.models.pixtral.convert_pixtral_weights_to_hf import MistralConverter
-
-    mistral_vocab = json.loads(data)
-
-    special_tokens = {}
-    vocab = {}
-
-    max_vocab = mistral_vocab["config"]["default_vocab_size"]
-    max_vocab -= len(mistral_vocab["special_tokens"])
-
-    for w in mistral_vocab["vocab"]:
-        r = w["rank"]
-        if r >= max_vocab:
-            continue
-
-        vocab[base64.b64decode(w["token_bytes"])] = r
-
-    for w in mistral_vocab["special_tokens"]:
-        if "token_bytes" in w:
-            special_tokens[base64.b64decode(w["token_bytes"])] = w["rank"]
-        else:
-            special_tokens[w["token_str"]] = w["rank"]
-
-    all_special = []
-    for v in special_tokens:
-        all_special.append(v)
-
-    special_tokens.update(vocab)
-    vocab = special_tokens
-    return {"tokenizer_object": MistralConverter(vocab=vocab, additional_special_tokens=all_special).converted(), "legacy": False}
-
-class MistralTokenizerClass:
-    @staticmethod
-    def from_pretrained(path, **kwargs):
-        return LlamaTokenizerFast(**kwargs)
-
-class Mistral3Tokenizer(sd1_clip.SDTokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        self.tekken_data = tokenizer_data.get("tekken_model", None)
-        super().__init__("", pad_with_end=False, embedding_size=5120, embedding_key='mistral3_24b', tokenizer_class=MistralTokenizerClass, has_end_token=False, pad_to_max_length=False, pad_token=11, max_length=99999999, min_length=1, pad_left=True, tokenizer_args=load_mistral_tokenizer(self.tekken_data), tokenizer_data=tokenizer_data)
-
-    def state_dict(self):
-        return {"tekken_model": self.tekken_data}
-
-class Flux2Tokenizer(sd1_clip.SD1Tokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="mistral3_24b", tokenizer=Mistral3Tokenizer)
-        self.llama_template = '[SYSTEM_PROMPT]You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\nattribution and actions without speculation.[/SYSTEM_PROMPT][INST]{}[/INST]'
-
-    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, **kwargs):
-        if llama_template is None:
-            llama_text = self.llama_template.format(text)
-        else:
-            llama_text = llama_template.format(text)
-
-        tokens = super().tokenize_with_weights(llama_text, return_word_ids=return_word_ids, disable_weights=True, **kwargs)
-        return tokens
-
-class Mistral3_24BModel(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer=[10, 20, 30], layer_idx=None, dtype=None, attention_mask=True, model_options={}):
-        textmodel_json_config = {}
-        num_layers = model_options.get("num_layers", None)
-        if num_layers is not None:
-            textmodel_json_config["num_hidden_layers"] = num_layers
-            if num_layers < 40:
-                textmodel_json_config["final_norm"] = False
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"start": 1, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Mistral3Small24B, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
-
-class Flux2TEModel(sd1_clip.SD1ClipModel):
-    def __init__(self, device="cpu", dtype=None, model_options={}, name="mistral3_24b", clip_model=Mistral3_24BModel):
-        super().__init__(device=device, dtype=dtype, name=name, clip_model=clip_model, model_options=model_options)
-
-    def encode_token_weights(self, token_weight_pairs):
-        out, pooled, extra = super().encode_token_weights(token_weight_pairs)
-
-        out = torch.stack((out[:, 0], out[:, 1], out[:, 2]), dim=1)
-        out = out.movedim(1, 2)
-        out = out.reshape(out.shape[0], out.shape[1], -1)
-        return out, pooled, extra
-
-def flux2_te(dtype_llama=None, llama_quantization_metadata=None, pruned=False):
-    class Flux2TEModel_(Flux2TEModel):
-        def __init__(self, device="cpu", dtype=None, model_options={}):
-            if dtype_llama is not None:
-                dtype = dtype_llama
-            if llama_quantization_metadata is not None:
-                model_options = model_options.copy()
-                model_options["quantization_metadata"] = llama_quantization_metadata
-            if pruned:
-                model_options = model_options.copy()
-                model_options["num_layers"] = 30
-            super().__init__(device=device, dtype=dtype, model_options=model_options)
-    return Flux2TEModel_

comfy/text_encoders/genmo.py

@@ -26,12 +26,12 @@ class MochiT5Tokenizer(sd1_clip.SD1Tokenizer):
         super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, clip_name="t5xxl", tokenizer=T5XXLTokenizer)
 
 
-def mochi_te(dtype_t5=None, t5_quantization_metadata=None):
+def mochi_te(dtype_t5=None, t5xxl_scaled_fp8=None):
     class MochiTEModel_(MochiT5XXL):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if t5_quantization_metadata is not None:
+            if t5xxl_scaled_fp8 is not None and "t5xxl_scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["t5xxl_quantization_metadata"] = t5_quantization_metadata
+                model_options["t5xxl_scaled_fp8"] = t5xxl_scaled_fp8
             if dtype is None:
                 dtype = dtype_t5
             super().__init__(device=device, dtype=dtype, model_options=model_options)

comfy/text_encoders/hidream.py

@@ -142,14 +142,14 @@ class HiDreamTEModel(torch.nn.Module):
             return self.llama.load_sd(sd)
 
 
-def hidream_clip(clip_l=True, clip_g=True, t5=True, llama=True, dtype_t5=None, dtype_llama=None, t5_quantization_metadata=None, llama_quantization_metadata=None):
+def hidream_clip(clip_l=True, clip_g=True, t5=True, llama=True, dtype_t5=None, dtype_llama=None, t5xxl_scaled_fp8=None, llama_scaled_fp8=None):
     class HiDreamTEModel_(HiDreamTEModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if t5_quantization_metadata is not None:
+            if t5xxl_scaled_fp8 is not None and "t5xxl_scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["t5xxl_quantization_metadata"] = t5_quantization_metadata
-            if llama_quantization_metadata is not None:
+                model_options["t5xxl_scaled_fp8"] = t5xxl_scaled_fp8
+            if llama_scaled_fp8 is not None and "llama_scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["llama_quantization_metadata"] = llama_quantization_metadata
+                model_options["llama_scaled_fp8"] = llama_scaled_fp8
             super().__init__(clip_l=clip_l, clip_g=clip_g, t5=t5, llama=llama, dtype_t5=dtype_t5, dtype_llama=dtype_llama, device=device, dtype=dtype, model_options=model_options)
     return HiDreamTEModel_

comfy/text_encoders/hunyuan_image.py

@@ -40,10 +40,10 @@ class HunyuanImageTokenizer(QwenImageTokenizer):
 
 class Qwen25_7BVLIModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="hidden", layer_idx=-3, dtype=None, attention_mask=True, model_options={}):
-        llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
-        if llama_quantization_metadata is not None:
+        llama_scaled_fp8 = model_options.get("qwen_scaled_fp8", None)
+        if llama_scaled_fp8 is not None:
             model_options = model_options.copy()
-            model_options["quantization_metadata"] = llama_quantization_metadata
+            model_options["scaled_fp8"] = llama_scaled_fp8
         super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Qwen25_7BVLI, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
 
@@ -91,12 +91,12 @@ class HunyuanImageTEModel(QwenImageTEModel):
         else:
             return super().load_sd(sd)
 
-def te(byt5=True, dtype_llama=None, llama_quantization_metadata=None):
+def te(byt5=True, dtype_llama=None, llama_scaled_fp8=None):
     class QwenImageTEModel_(HunyuanImageTEModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if llama_quantization_metadata is not None:
+            if llama_scaled_fp8 is not None and "scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["llama_quantization_metadata"] = llama_quantization_metadata
+                model_options["qwen_scaled_fp8"] = llama_scaled_fp8
             if dtype_llama is not None:
                 dtype = dtype_llama
             super().__init__(byt5=byt5, device=device, dtype=dtype, model_options=model_options)

comfy/text_encoders/hunyuan_video.py

@@ -1,12 +1,11 @@
 from comfy import sd1_clip
 import comfy.model_management
 import comfy.text_encoders.llama
-from .hunyuan_image import HunyuanImageTokenizer
 from transformers import LlamaTokenizerFast
 import torch
 import os
 import numbers
-import comfy.utils
+
 
 def llama_detect(state_dict, prefix=""):
     out = {}
@@ -14,9 +13,9 @@ def llama_detect(state_dict, prefix=""):
     if t5_key in state_dict:
         out["dtype_llama"] = state_dict[t5_key].dtype
 
-    quant = comfy.utils.detect_layer_quantization(state_dict, prefix)
-    if quant is not None:
-        out["llama_quantization_metadata"] = quant
+    scaled_fp8_key = "{}scaled_fp8".format(prefix)
+    if scaled_fp8_key in state_dict:
+        out["llama_scaled_fp8"] = state_dict[scaled_fp8_key].dtype
 
     return out
 
@@ -28,10 +27,10 @@ class LLAMA3Tokenizer(sd1_clip.SDTokenizer):
 
 class LLAMAModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="hidden", layer_idx=-3, dtype=None, attention_mask=True, model_options={}, special_tokens={"start": 128000, "pad": 128258}):
-        llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
-        if llama_quantization_metadata is not None:
+        llama_scaled_fp8 = model_options.get("llama_scaled_fp8", None)
+        if llama_scaled_fp8 is not None:
             model_options = model_options.copy()
-            model_options["quantization_metadata"] = llama_quantization_metadata
+            model_options["scaled_fp8"] = llama_scaled_fp8
 
         textmodel_json_config = {}
         vocab_size = model_options.get("vocab_size", None)
@@ -74,14 +73,6 @@ class HunyuanVideoTokenizer:
         return {}
 
 
-class HunyuanVideo15Tokenizer(HunyuanImageTokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
-        self.llama_template = "<|im_start|>system\nYou are a helpful assistant. Describe the video by detailing the following aspects:\n1. The main content and theme of the video.\n2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.\n3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.\n4. background environment, light, style and atmosphere.\n5. camera angles, movements, and transitions used in the video.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
-
-    def tokenize_with_weights(self, text:str, return_word_ids=False, **kwargs):
-        return super().tokenize_with_weights(text, return_word_ids, prevent_empty_text=True, **kwargs)
-
 class HunyuanVideoClipModel(torch.nn.Module):
     def __init__(self, dtype_llama=None, device="cpu", dtype=None, model_options={}):
         super().__init__()
@@ -158,11 +149,11 @@ class HunyuanVideoClipModel(torch.nn.Module):
             return self.llama.load_sd(sd)
 
 
-def hunyuan_video_clip(dtype_llama=None, llama_quantization_metadata=None):
+def hunyuan_video_clip(dtype_llama=None, llama_scaled_fp8=None):
     class HunyuanVideoClipModel_(HunyuanVideoClipModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if llama_quantization_metadata is not None:
+            if llama_scaled_fp8 is not None and "llama_scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["llama_quantization_metadata"] = llama_quantization_metadata
+                model_options["llama_scaled_fp8"] = llama_scaled_fp8
             super().__init__(dtype_llama=dtype_llama, device=device, dtype=dtype, model_options=model_options)
     return HunyuanVideoClipModel_

comfy/text_encoders/jina_clip_2.py

@@ -1,219 +0,0 @@
-# Jina CLIP v2 and Jina Embeddings v3 both use their modified XLM-RoBERTa architecture. Reference implementation:
-# Jina CLIP v2 (both text and vision): https://huggingface.co/jinaai/jina-clip-implementation/blob/39e6a55ae971b59bea6e44675d237c99762e7ee2/modeling_clip.py
-# Jina XLM-RoBERTa (text only): http://huggingface.co/jinaai/xlm-roberta-flash-implementation/blob/2b6bc3f30750b3a9648fe9b63448c09920efe9be/modeling_xlm_roberta.py
-
-from dataclasses import dataclass
-
-import torch
-from torch import nn as nn
-from torch.nn import functional as F
-
-import comfy.model_management
-import comfy.ops
-from comfy import sd1_clip
-from .spiece_tokenizer import SPieceTokenizer
-
-class JinaClip2Tokenizer(sd1_clip.SDTokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        tokenizer = tokenizer_data.get("spiece_model", None)
-        # The official NewBie uses max_length=8000, but Jina Embeddings v3 actually supports 8192
-        super().__init__(tokenizer, pad_with_end=False, embedding_size=1024, embedding_key='jina_clip_2', tokenizer_class=SPieceTokenizer, has_start_token=True, has_end_token=True, pad_to_max_length=False, max_length=8192, min_length=1, pad_token=1, end_token=2, tokenizer_args={"add_bos": True, "add_eos": True}, tokenizer_data=tokenizer_data)
-
-    def state_dict(self):
-        return {"spiece_model": self.tokenizer.serialize_model()}
-
-class JinaClip2TokenizerWrapper(sd1_clip.SD1Tokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, tokenizer=JinaClip2Tokenizer, name="jina_clip_2")
-
-# https://huggingface.co/jinaai/jina-embeddings-v3/blob/343dbf534c76fe845f304fa5c2d1fd87e1e78918/config.json
-@dataclass
-class XLMRobertaConfig:
-    vocab_size: int = 250002
-    type_vocab_size: int = 1
-    hidden_size: int = 1024
-    num_hidden_layers: int = 24
-    num_attention_heads: int = 16
-    rotary_emb_base: float = 20000.0
-    intermediate_size: int = 4096
-    hidden_act: str = "gelu"
-    hidden_dropout_prob: float = 0.1
-    attention_probs_dropout_prob: float = 0.1
-    layer_norm_eps: float = 1e-05
-    bos_token_id: int = 0
-    eos_token_id: int = 2
-    pad_token_id: int = 1
-
-class XLMRobertaEmbeddings(nn.Module):
-    def __init__(self, config, device=None, dtype=None, ops=None):
-        super().__init__()
-        embed_dim = config.hidden_size
-        self.word_embeddings = ops.Embedding(config.vocab_size, embed_dim, padding_idx=config.pad_token_id, device=device, dtype=dtype)
-        self.token_type_embeddings = ops.Embedding(config.type_vocab_size, embed_dim, device=device, dtype=dtype)
-
-    def forward(self, input_ids=None, embeddings=None):
-        if input_ids is not None and embeddings is None:
-            embeddings = self.word_embeddings(input_ids)
-
-        if embeddings is not None:
-            token_type_ids = torch.zeros(embeddings.shape[1], device=embeddings.device, dtype=torch.int32)
-            token_type_embeddings = self.token_type_embeddings(token_type_ids)
-            embeddings = embeddings + token_type_embeddings
-        return embeddings
-
-class RotaryEmbedding(nn.Module):
-    def __init__(self, dim, base, device=None):
-        super().__init__()
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-        self._seq_len_cached = 0
-        self._cos_cached = None
-        self._sin_cached = None
-
-    def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
-        if seqlen > self._seq_len_cached or self._cos_cached is None or self._cos_cached.device != device or self._cos_cached.dtype != dtype:
-            self._seq_len_cached = seqlen
-            t = torch.arange(seqlen, device=device, dtype=torch.float32)
-            freqs = torch.outer(t, self.inv_freq.to(device=t.device))
-            emb = torch.cat((freqs, freqs), dim=-1)
-            self._cos_cached = emb.cos().to(dtype)
-            self._sin_cached = emb.sin().to(dtype)
-
-    def forward(self, q, k):
-        batch, seqlen, heads, head_dim = q.shape
-        self._update_cos_sin_cache(seqlen, device=q.device, dtype=q.dtype)
-
-        cos = self._cos_cached[:seqlen].view(1, seqlen, 1, head_dim)
-        sin = self._sin_cached[:seqlen].view(1, seqlen, 1, head_dim)
-
-        def rotate_half(x):
-            size = x.shape[-1] // 2
-            x1, x2 = x[..., :size], x[..., size:]
-            return torch.cat((-x2, x1), dim=-1)
-
-        q_embed = (q * cos) + (rotate_half(q) * sin)
-        k_embed = (k * cos) + (rotate_half(k) * sin)
-        return q_embed, k_embed
-
-class MHA(nn.Module):
-    def __init__(self, config, device=None, dtype=None, ops=None):
-        super().__init__()
-        embed_dim = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.head_dim = embed_dim // config.num_attention_heads
-
-        self.rotary_emb = RotaryEmbedding(self.head_dim, config.rotary_emb_base, device=device)
-        self.Wqkv = ops.Linear(embed_dim, 3 * embed_dim, device=device, dtype=dtype)
-        self.out_proj = ops.Linear(embed_dim, embed_dim, device=device, dtype=dtype)
-
-    def forward(self, x, mask=None, optimized_attention=None):
-        qkv = self.Wqkv(x)
-        batch_size, seq_len, _ = qkv.shape
-        qkv = qkv.view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
-        q, k, v = qkv.unbind(2)
-
-        q, k = self.rotary_emb(q, k)
-
-        # NHD -> HND
-        q = q.transpose(1, 2)
-        k = k.transpose(1, 2)
-        v = v.transpose(1, 2)
-
-        out = optimized_attention(q, k, v, heads=self.num_heads, mask=mask, skip_reshape=True)
-        return self.out_proj(out)
-
-class MLP(nn.Module):
-    def __init__(self, config, device=None, dtype=None, ops=None):
-        super().__init__()
-        self.fc1 = ops.Linear(config.hidden_size, config.intermediate_size, device=device, dtype=dtype)
-        self.activation = F.gelu
-        self.fc2 = ops.Linear(config.intermediate_size, config.hidden_size, device=device, dtype=dtype)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.activation(x)
-        x = self.fc2(x)
-        return x
-
-class Block(nn.Module):
-    def __init__(self, config, device=None, dtype=None, ops=None):
-        super().__init__()
-        self.mixer = MHA(config, device=device, dtype=dtype, ops=ops)
-        self.dropout1 = nn.Dropout(config.hidden_dropout_prob)
-        self.norm1 = ops.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device, dtype=dtype)
-        self.mlp = MLP(config, device=device, dtype=dtype, ops=ops)
-        self.dropout2 = nn.Dropout(config.hidden_dropout_prob)
-        self.norm2 = ops.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device, dtype=dtype)
-
-    def forward(self, hidden_states, mask=None, optimized_attention=None):
-        mixer_out = self.mixer(hidden_states, mask=mask, optimized_attention=optimized_attention)
-        hidden_states = self.norm1(self.dropout1(mixer_out) + hidden_states)
-        mlp_out = self.mlp(hidden_states)
-        hidden_states = self.norm2(self.dropout2(mlp_out) + hidden_states)
-        return hidden_states
-
-class XLMRobertaEncoder(nn.Module):
-    def __init__(self, config, device=None, dtype=None, ops=None):
-        super().__init__()
-        self.layers = nn.ModuleList([Block(config, device=device, dtype=dtype, ops=ops) for _ in range(config.num_hidden_layers)])
-
-    def forward(self, hidden_states, attention_mask=None):
-        optimized_attention = comfy.ldm.modules.attention.optimized_attention_for_device(hidden_states.device, mask=attention_mask is not None, small_input=True)
-        for layer in self.layers:
-            hidden_states = layer(hidden_states, mask=attention_mask, optimized_attention=optimized_attention)
-        return hidden_states
-
-class XLMRobertaModel_(nn.Module):
-    def __init__(self, config, device=None, dtype=None, ops=None):
-        super().__init__()
-        self.embeddings = XLMRobertaEmbeddings(config, device=device, dtype=dtype, ops=ops)
-        self.emb_ln = ops.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, device=device, dtype=dtype)
-        self.emb_drop = nn.Dropout(config.hidden_dropout_prob)
-        self.encoder = XLMRobertaEncoder(config, device=device, dtype=dtype, ops=ops)
-
-    def forward(self, input_ids, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, embeds_info=[]):
-        x = self.embeddings(input_ids=input_ids, embeddings=embeds)
-        x = self.emb_ln(x)
-        x = self.emb_drop(x)
-
-        mask = None
-        if attention_mask is not None:
-            mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, 1, attention_mask.shape[-1]))
-            mask = mask.masked_fill(mask.to(torch.bool), -torch.finfo(x.dtype).max)
-
-        sequence_output = self.encoder(x, attention_mask=mask)
-
-        # Mean pool, see https://huggingface.co/jinaai/jina-clip-implementation/blob/39e6a55ae971b59bea6e44675d237c99762e7ee2/hf_model.py
-        pooled_output = None
-        if attention_mask is None:
-            pooled_output = sequence_output.mean(dim=1)
-        else:
-            attention_mask = attention_mask.to(sequence_output.dtype)
-            pooled_output = (sequence_output * attention_mask.unsqueeze(-1)).sum(dim=1) / attention_mask.sum(dim=-1, keepdim=True)
-
-        # Intermediate output is not yet implemented, use None for placeholder
-        return sequence_output, None, pooled_output
-
-class XLMRobertaModel(nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        self.config = XLMRobertaConfig(**config_dict)
-        self.model = XLMRobertaModel_(self.config, device=device, dtype=dtype, ops=operations)
-        self.num_layers = self.config.num_hidden_layers
-
-    def get_input_embeddings(self):
-        return self.model.embeddings.word_embeddings
-
-    def set_input_embeddings(self, embeddings):
-        self.model.embeddings.word_embeddings = embeddings
-
-    def forward(self, *args, **kwargs):
-        return self.model(*args, **kwargs)
-
-class JinaClip2TextModel(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", dtype=None, model_options={}):
-        super().__init__(device=device, dtype=dtype, textmodel_json_config={}, model_class=XLMRobertaModel, special_tokens={"start": 0, "end": 2, "pad": 1}, enable_attention_masks=True, return_attention_masks=True, model_options=model_options)
-
-class JinaClip2TextModelWrapper(sd1_clip.SD1ClipModel):
-    def __init__(self, device="cpu", dtype=None, model_options={}):
-        super().__init__(device=device, dtype=dtype, clip_model=JinaClip2TextModel, name="jina_clip_2", model_options=model_options)

comfy/text_encoders/kandinsky5.py

@@ -1,68 +0,0 @@
-from comfy import sd1_clip
-from .qwen_image import QwenImageTokenizer, QwenImageTEModel
-from .llama import Qwen25_7BVLI
-
-
-class Kandinsky5Tokenizer(QwenImageTokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
-        self.llama_template = "<|im_start|>system\nYou are a prompt engineer. Describe the video in detail.\nDescribe how the camera moves or shakes, describe the zoom and view angle, whether it follows the objects.\nDescribe the location of the video, main characters or objects and their action.\nDescribe the dynamism of the video and presented actions.\nName the visual style of the video: whether it is a professional footage, user generated content, some kind of animation, video game or screen content.\nDescribe the visual effects, postprocessing and transitions if they are presented in the video.\nPay attention to the order of key actions shown in the scene.<|im_end|>\n<|im_start|>user\n{}<|im_end|>"
-        self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
-
-    def tokenize_with_weights(self, text:str, return_word_ids=False, **kwargs):
-        out = super().tokenize_with_weights(text, return_word_ids, **kwargs)
-        out["l"] = self.clip_l.tokenize_with_weights(text, return_word_ids, **kwargs)
-
-        return out
-
-
-class Kandinsky5TokenizerImage(Kandinsky5Tokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
-        self.llama_template = "<|im_start|>system\nYou are a promt engineer. Describe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>"
-
-
-class Qwen25_7BVLIModel(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="hidden", layer_idx=-1, dtype=None, attention_mask=True, model_options={}):
-        llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
-        if llama_quantization_metadata is not None:
-            model_options = model_options.copy()
-            model_options["quantization_metadata"] = llama_quantization_metadata
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False, model_class=Qwen25_7BVLI, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
-
-
-class Kandinsky5TEModel(QwenImageTEModel):
-    def __init__(self, device="cpu", dtype=None, model_options={}):
-        super(QwenImageTEModel, self).__init__(device=device, dtype=dtype, name="qwen25_7b", clip_model=Qwen25_7BVLIModel, model_options=model_options)
-        self.clip_l = sd1_clip.SDClipModel(device=device, dtype=dtype, return_projected_pooled=False, model_options=model_options)
-
-    def encode_token_weights(self, token_weight_pairs):
-        cond, p, extra = super().encode_token_weights(token_weight_pairs, template_end=-1)
-        l_out, l_pooled = self.clip_l.encode_token_weights(token_weight_pairs["l"])
-
-        return cond, l_pooled, extra
-
-    def set_clip_options(self, options):
-        super().set_clip_options(options)
-        self.clip_l.set_clip_options(options)
-
-    def reset_clip_options(self):
-        super().reset_clip_options()
-        self.clip_l.reset_clip_options()
-
-    def load_sd(self, sd):
-        if "text_model.encoder.layers.1.mlp.fc1.weight" in sd:
-            return self.clip_l.load_sd(sd)
-        else:
-            return super().load_sd(sd)
-
-def te(dtype_llama=None, llama_quantization_metadata=None):
-    class Kandinsky5TEModel_(Kandinsky5TEModel):
-        def __init__(self, device="cpu", dtype=None, model_options={}):
-            if llama_quantization_metadata is not None:
-                model_options = model_options.copy()
-                model_options["llama_quantization_metadata"] = llama_quantization_metadata
-            if dtype_llama is not None:
-                dtype = dtype_llama
-            super().__init__(device=device, dtype=dtype, model_options=model_options)
-    return Kandinsky5TEModel_

comfy/text_encoders/llama.py

@@ -3,6 +3,7 @@ import torch.nn as nn
 from dataclasses import dataclass
 from typing import Optional, Any
 import math
+import logging
 
 from comfy.ldm.modules.attention import optimized_attention_for_device
 import comfy.model_management
@@ -31,29 +32,6 @@ class Llama2Config:
     q_norm = None
     k_norm = None
     rope_scale = None
-    final_norm: bool = True
-
-@dataclass
-class Mistral3Small24BConfig:
-    vocab_size: int = 131072
-    hidden_size: int = 5120
-    intermediate_size: int = 32768
-    num_hidden_layers: int = 40
-    num_attention_heads: int = 32
-    num_key_value_heads: int = 8
-    max_position_embeddings: int = 8192
-    rms_norm_eps: float = 1e-5
-    rope_theta: float = 1000000000.0
-    transformer_type: str = "llama"
-    head_dim = 128
-    rms_norm_add = False
-    mlp_activation = "silu"
-    qkv_bias = False
-    rope_dims = None
-    q_norm = None
-    k_norm = None
-    rope_scale = None
-    final_norm: bool = True
 
 @dataclass
 class Qwen25_3BConfig:
@@ -75,51 +53,6 @@ class Qwen25_3BConfig:
     q_norm = None
     k_norm = None
     rope_scale = None
-    final_norm: bool = True
-
-@dataclass
-class Qwen3_4BConfig:
-    vocab_size: int = 151936
-    hidden_size: int = 2560
-    intermediate_size: int = 9728
-    num_hidden_layers: int = 36
-    num_attention_heads: int = 32
-    num_key_value_heads: int = 8
-    max_position_embeddings: int = 40960
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 1000000.0
-    transformer_type: str = "llama"
-    head_dim = 128
-    rms_norm_add = False
-    mlp_activation = "silu"
-    qkv_bias = False
-    rope_dims = None
-    q_norm = "gemma3"
-    k_norm = "gemma3"
-    rope_scale = None
-    final_norm: bool = True
-
-@dataclass
-class Ovis25_2BConfig:
-    vocab_size: int = 151936
-    hidden_size: int = 2048
-    intermediate_size: int = 6144
-    num_hidden_layers: int = 28
-    num_attention_heads: int = 16
-    num_key_value_heads: int = 8
-    max_position_embeddings: int = 40960
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 1000000.0
-    transformer_type: str = "llama"
-    head_dim = 128
-    rms_norm_add = False
-    mlp_activation = "silu"
-    qkv_bias = False
-    rope_dims = None
-    q_norm = "gemma3"
-    k_norm = "gemma3"
-    rope_scale = None
-    final_norm: bool = True
 
 @dataclass
 class Qwen25_7BVLI_Config:
@@ -141,7 +74,6 @@ class Qwen25_7BVLI_Config:
     q_norm = None
     k_norm = None
     rope_scale = None
-    final_norm: bool = True
 
 @dataclass
 class Gemma2_2B_Config:
@@ -164,7 +96,6 @@ class Gemma2_2B_Config:
     k_norm = None
     sliding_attention = None
     rope_scale = None
-    final_norm: bool = True
 
 @dataclass
 class Gemma3_4B_Config:
@@ -176,7 +107,7 @@ class Gemma3_4B_Config:
     num_key_value_heads: int = 4
     max_position_embeddings: int = 131072
     rms_norm_eps: float = 1e-6
-    rope_theta = [1000000.0, 10000.0]
+    rope_theta = [10000.0, 1000000.0]
     transformer_type: str = "gemma3"
     head_dim = 256
     rms_norm_add = True
@@ -185,9 +116,8 @@ class Gemma3_4B_Config:
     rope_dims = None
     q_norm = "gemma3"
     k_norm = "gemma3"
-    sliding_attention = [1024, 1024, 1024, 1024, 1024, False]
-    rope_scale = [8.0, 1.0]
-    final_norm: bool = True
+    sliding_attention = [False, False, False, False, False, 1024]
+    rope_scale = [1.0, 8.0]
 
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-5, add=False, device=None, dtype=None):
@@ -369,7 +299,7 @@ class TransformerBlockGemma2(nn.Module):
         self.pre_feedforward_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
         self.post_feedforward_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
 
-        if config.sliding_attention is not None:
+        if config.sliding_attention is not None:  # TODO: implement. (Not that necessary since models are trained on less than 1024 tokens)
             self.sliding_attention = config.sliding_attention[index % len(config.sliding_attention)]
         else:
             self.sliding_attention = False
@@ -386,12 +316,7 @@ class TransformerBlockGemma2(nn.Module):
         if self.transformer_type == 'gemma3':
             if self.sliding_attention:
                 if x.shape[1] > self.sliding_attention:
-                    sliding_mask = torch.full((x.shape[1], x.shape[1]), float("-inf"), device=x.device, dtype=x.dtype)
-                    sliding_mask.tril_(diagonal=-self.sliding_attention)
-                    if attention_mask is not None:
-                        attention_mask = attention_mask + sliding_mask
-                    else:
-                        attention_mask = sliding_mask
+                    logging.warning("Warning: sliding attention not implemented, results may be incorrect")
                 freqs_cis = freqs_cis[1]
             else:
                 freqs_cis = freqs_cis[0]
@@ -441,12 +366,7 @@ class Llama2_(nn.Module):
             transformer(config, index=i, device=device, dtype=dtype, ops=ops)
             for i in range(config.num_hidden_layers)
         ])
-
-        if config.final_norm:
-            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
-        else:
-            self.norm = None
-
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, add=config.rms_norm_add, device=device, dtype=dtype)
         # self.lm_head = ops.Linear(config.hidden_size, config.vocab_size, bias=False, device=device, dtype=dtype)
 
     def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, position_ids=None, embeds_info=[]):
@@ -482,12 +402,8 @@ class Llama2_(nn.Module):
 
         intermediate = None
         all_intermediate = None
-        only_layers = None
         if intermediate_output is not None:
-            if isinstance(intermediate_output, list):
-                all_intermediate = []
-                only_layers = set(intermediate_output)
-            elif intermediate_output == "all":
+            if intermediate_output == "all":
                 all_intermediate = []
                 intermediate_output = None
             elif intermediate_output < 0:
@@ -495,8 +411,7 @@ class Llama2_(nn.Module):
 
         for i, layer in enumerate(self.layers):
             if all_intermediate is not None:
-                if only_layers is None or (i in only_layers):
-                    all_intermediate.append(x.unsqueeze(1).clone())
+                all_intermediate.append(x.unsqueeze(1).clone())
             x = layer(
                 x=x,
                 attention_mask=mask,
@@ -506,17 +421,14 @@ class Llama2_(nn.Module):
             if i == intermediate_output:
                 intermediate = x.clone()
 
-        if self.norm is not None:
-            x = self.norm(x)
-
+        x = self.norm(x)
         if all_intermediate is not None:
-            if only_layers is None or ((i + 1) in only_layers):
-                all_intermediate.append(x.unsqueeze(1).clone())
+            all_intermediate.append(x.unsqueeze(1).clone())
 
         if all_intermediate is not None:
             intermediate = torch.cat(all_intermediate, dim=1)
 
-        if intermediate is not None and final_layer_norm_intermediate and self.norm is not None:
+        if intermediate is not None and final_layer_norm_intermediate:
             intermediate = self.norm(intermediate)
 
         return x, intermediate
@@ -541,15 +453,6 @@ class Llama2(BaseLlama, torch.nn.Module):
         self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
         self.dtype = dtype
 
-class Mistral3Small24B(BaseLlama, torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        config = Mistral3Small24BConfig(**config_dict)
-        self.num_layers = config.num_hidden_layers
-
-        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
-        self.dtype = dtype
-
 class Qwen25_3B(BaseLlama, torch.nn.Module):
     def __init__(self, config_dict, dtype, device, operations):
         super().__init__()
@@ -559,24 +462,6 @@ class Qwen25_3B(BaseLlama, torch.nn.Module):
         self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
         self.dtype = dtype
 
-class Qwen3_4B(BaseLlama, torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        config = Qwen3_4BConfig(**config_dict)
-        self.num_layers = config.num_hidden_layers
-
-        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
-        self.dtype = dtype
-
-class Ovis25_2B(BaseLlama, torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        config = Ovis25_2BConfig(**config_dict)
-        self.num_layers = config.num_hidden_layers
-
-        self.model = Llama2_(config, device=device, dtype=dtype, ops=operations)
-        self.dtype = dtype
-
 class Qwen25_7BVLI(BaseLlama, torch.nn.Module):
     def __init__(self, config_dict, dtype, device, operations):
         super().__init__()

comfy/text_encoders/lumina2.py

@@ -14,7 +14,7 @@ class Gemma2BTokenizer(sd1_clip.SDTokenizer):
 class Gemma3_4BTokenizer(sd1_clip.SDTokenizer):
     def __init__(self, embedding_directory=None, tokenizer_data={}):
         tokenizer = tokenizer_data.get("spiece_model", None)
-        super().__init__(tokenizer, pad_with_end=False, embedding_size=2560, embedding_key='gemma3_4b', tokenizer_class=SPieceTokenizer, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, tokenizer_args={"add_bos": True, "add_eos": False}, disable_weights=True, tokenizer_data=tokenizer_data)
+        super().__init__(tokenizer, pad_with_end=False, embedding_size=2560, embedding_key='gemma3_4b', tokenizer_class=SPieceTokenizer, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, tokenizer_args={"add_bos": True, "add_eos": False}, tokenizer_data=tokenizer_data)
 
     def state_dict(self):
         return {"spiece_model": self.tokenizer.serialize_model()}
@@ -33,11 +33,6 @@ class Gemma2_2BModel(sd1_clip.SDClipModel):
 
 class Gemma3_4BModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="hidden", layer_idx=-2, dtype=None, attention_mask=True, model_options={}):
-        llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
-        if llama_quantization_metadata is not None:
-            model_options = model_options.copy()
-            model_options["quantization_metadata"] = llama_quantization_metadata
-
         super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Gemma3_4B, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
 class LuminaModel(sd1_clip.SD1ClipModel):
@@ -45,7 +40,7 @@ class LuminaModel(sd1_clip.SD1ClipModel):
         super().__init__(device=device, dtype=dtype, name=name, clip_model=clip_model, model_options=model_options)
 
 
-def te(dtype_llama=None, llama_quantization_metadata=None, model_type="gemma2_2b"):
+def te(dtype_llama=None, llama_scaled_fp8=None, model_type="gemma2_2b"):
     if model_type == "gemma2_2b":
         model = Gemma2_2BModel
     elif model_type == "gemma3_4b":
@@ -53,9 +48,9 @@ def te(dtype_llama=None, llama_quantization_metadata=None, model_type="gemma2_2b
 
     class LuminaTEModel_(LuminaModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if llama_quantization_metadata is not None:
+            if llama_scaled_fp8 is not None and "scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["quantization_metadata"] = llama_quantization_metadata
+                model_options["scaled_fp8"] = llama_scaled_fp8
             if dtype_llama is not None:
                 dtype = dtype_llama
             super().__init__(device=device, dtype=dtype, name=model_type, model_options=model_options, clip_model=model)

comfy/text_encoders/newbie.py

@@ -1,62 +0,0 @@
-import torch
-
-import comfy.model_management
-import comfy.text_encoders.jina_clip_2
-import comfy.text_encoders.lumina2
-
-class NewBieTokenizer:
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        self.gemma = comfy.text_encoders.lumina2.Gemma3_4BTokenizer(embedding_directory=embedding_directory, tokenizer_data={"spiece_model": tokenizer_data["gemma_spiece_model"]})
-        self.jina = comfy.text_encoders.jina_clip_2.JinaClip2Tokenizer(embedding_directory=embedding_directory, tokenizer_data={"spiece_model": tokenizer_data["jina_spiece_model"]})
-
-    def tokenize_with_weights(self, text:str, return_word_ids=False, **kwargs):
-        out = {}
-        out["gemma"] = self.gemma.tokenize_with_weights(text, return_word_ids, **kwargs)
-        out["jina"] = self.jina.tokenize_with_weights(text, return_word_ids, **kwargs)
-        return out
-
-    def untokenize(self, token_weight_pair):
-        raise NotImplementedError
-
-    def state_dict(self):
-        return {}
-
-class NewBieTEModel(torch.nn.Module):
-    def __init__(self, dtype_gemma=None, device="cpu", dtype=None, model_options={}):
-        super().__init__()
-        dtype_gemma = comfy.model_management.pick_weight_dtype(dtype_gemma, dtype, device)
-        self.gemma = comfy.text_encoders.lumina2.Gemma3_4BModel(device=device, dtype=dtype_gemma, model_options=model_options)
-        self.jina = comfy.text_encoders.jina_clip_2.JinaClip2TextModel(device=device, dtype=dtype, model_options=model_options)
-        self.dtypes = {dtype, dtype_gemma}
-
-    def set_clip_options(self, options):
-        self.gemma.set_clip_options(options)
-        self.jina.set_clip_options(options)
-
-    def reset_clip_options(self):
-        self.gemma.reset_clip_options()
-        self.jina.reset_clip_options()
-
-    def encode_token_weights(self, token_weight_pairs):
-        token_weight_pairs_gemma = token_weight_pairs["gemma"]
-        token_weight_pairs_jina = token_weight_pairs["jina"]
-
-        gemma_out, gemma_pooled, gemma_extra = self.gemma.encode_token_weights(token_weight_pairs_gemma)
-        jina_out, jina_pooled, jina_extra = self.jina.encode_token_weights(token_weight_pairs_jina)
-
-        return gemma_out, jina_pooled, gemma_extra
-
-    def load_sd(self, sd):
-        if "model.layers.0.self_attn.q_norm.weight" in sd:
-            return self.gemma.load_sd(sd)
-        else:
-            return self.jina.load_sd(sd)
-
-def te(dtype_llama=None, llama_quantization_metadata=None):
-    class NewBieTEModel_(NewBieTEModel):
-        def __init__(self, device="cpu", dtype=None, model_options={}):
-            if llama_quantization_metadata is not None:
-                model_options = model_options.copy()
-                model_options["llama_quantization_metadata"] = llama_quantization_metadata
-            super().__init__(dtype_gemma=dtype_llama, device=device, dtype=dtype, model_options=model_options)
-    return NewBieTEModel_

comfy/text_encoders/omnigen2.py

@@ -32,12 +32,12 @@ class Omnigen2Model(sd1_clip.SD1ClipModel):
         super().__init__(device=device, dtype=dtype, name="qwen25_3b", clip_model=Qwen25_3BModel, model_options=model_options)
 
 
-def te(dtype_llama=None, llama_quantization_metadata=None):
+def te(dtype_llama=None, llama_scaled_fp8=None):
     class Omnigen2TEModel_(Omnigen2Model):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if llama_quantization_metadata is not None:
+            if llama_scaled_fp8 is not None and "scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["quantization_metadata"] = llama_quantization_metadata
+                model_options["scaled_fp8"] = llama_scaled_fp8
             if dtype_llama is not None:
                 dtype = dtype_llama
             super().__init__(device=device, dtype=dtype, model_options=model_options)

comfy/text_encoders/ovis.py

@@ -1,66 +0,0 @@
-from transformers import Qwen2Tokenizer
-import comfy.text_encoders.llama
-from comfy import sd1_clip
-import os
-import torch
-import numbers
-
-class Qwen3Tokenizer(sd1_clip.SDTokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "qwen25_tokenizer")
-        super().__init__(tokenizer_path, pad_with_end=False, embedding_size=2048, embedding_key='qwen3_2b', tokenizer_class=Qwen2Tokenizer, has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=284, pad_token=151643, tokenizer_data=tokenizer_data)
-
-
-class OvisTokenizer(sd1_clip.SD1Tokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="qwen3_2b", tokenizer=Qwen3Tokenizer)
-        self.llama_template = "<|im_start|>user\nDescribe the image by detailing the color, quantity, text, shape, size, texture, spatial relationships of the objects and background: {}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
-
-    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, **kwargs):
-        if llama_template is None:
-            llama_text = self.llama_template.format(text)
-        else:
-            llama_text = llama_template.format(text)
-
-        tokens = super().tokenize_with_weights(llama_text, return_word_ids=return_word_ids, disable_weights=True, **kwargs)
-        return tokens
-
-class Ovis25_2BModel(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Ovis25_2B, enable_attention_masks=attention_mask, return_attention_masks=False, zero_out_masked=True, model_options=model_options)
-
-
-class OvisTEModel(sd1_clip.SD1ClipModel):
-    def __init__(self, device="cpu", dtype=None, model_options={}):
-        super().__init__(device=device, dtype=dtype, name="qwen3_2b", clip_model=Ovis25_2BModel, model_options=model_options)
-
-    def encode_token_weights(self, token_weight_pairs, template_end=-1):
-        out, pooled = super().encode_token_weights(token_weight_pairs)
-        tok_pairs = token_weight_pairs["qwen3_2b"][0]
-        count_im_start = 0
-        if template_end == -1:
-            for i, v in enumerate(tok_pairs):
-                elem = v[0]
-                if not torch.is_tensor(elem):
-                    if isinstance(elem, numbers.Integral):
-                        if elem == 4004 and count_im_start < 1:
-                            template_end = i
-                            count_im_start += 1
-
-            if out.shape[1] > (template_end + 1):
-                if tok_pairs[template_end + 1][0] == 25:
-                    template_end += 1
-
-        out = out[:, template_end:]
-        return out, pooled, {}
-
-
-def te(dtype_llama=None, llama_quantization_metadata=None):
-    class OvisTEModel_(OvisTEModel):
-        def __init__(self, device="cpu", dtype=None, model_options={}):
-            if dtype_llama is not None:
-                dtype = dtype_llama
-            if llama_quantization_metadata is not None:
-                model_options["quantization_metadata"] = llama_quantization_metadata
-            super().__init__(device=device, dtype=dtype, model_options=model_options)
-    return OvisTEModel_

comfy/text_encoders/pixart_t5.py

@@ -30,12 +30,12 @@ class PixArtTokenizer(sd1_clip.SD1Tokenizer):
     def __init__(self, embedding_directory=None, tokenizer_data={}):
         super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, clip_name="t5xxl", tokenizer=T5XXLTokenizer)
 
-def pixart_te(dtype_t5=None, t5_quantization_metadata=None):
+def pixart_te(dtype_t5=None, t5xxl_scaled_fp8=None):
     class PixArtTEModel_(PixArtT5XXL):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if t5_quantization_metadata is not None:
+            if t5xxl_scaled_fp8 is not None and "t5xxl_scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["t5xxl_quantization_metadata"] = t5_quantization_metadata
+                model_options["t5xxl_scaled_fp8"] = t5xxl_scaled_fp8
             if dtype is None:
                 dtype = dtype_t5
             super().__init__(device=device, dtype=dtype, model_options=model_options)

comfy/text_encoders/qwen25_tokenizer/tokenizer_config.json

@@ -179,36 +179,36 @@
       "special": false
     },
     "151665": {
-      "content": "<tool_response>",
+      "content": "<|img|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
       "single_word": false,
-      "special": false
+      "special": true
     },
     "151666": {
-      "content": "</tool_response>",
+      "content": "<|endofimg|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
       "single_word": false,
-      "special": false
+      "special": true
     },
     "151667": {
-      "content": "<think>",
+      "content": "<|meta|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
       "single_word": false,
-      "special": false
+      "special": true
     },
     "151668": {
-      "content": "</think>",
+      "content": "<|endofmeta|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
       "single_word": false,
-      "special": false
+      "special": true
     }
   },
   "additional_special_tokens": [

comfy/text_encoders/qwen_image.py

@@ -17,14 +17,12 @@ class QwenImageTokenizer(sd1_clip.SD1Tokenizer):
         self.llama_template = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
         self.llama_template_images = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"
 
-    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, **kwargs):
+    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], **kwargs):
         skip_template = False
         if text.startswith('<|im_start|>'):
             skip_template = True
         if text.startswith('<|start_header_id|>'):
             skip_template = True
-        if prevent_empty_text and text == '':
-            text = ' '
 
         if skip_template:
             llama_text = text
@@ -85,12 +83,12 @@ class QwenImageTEModel(sd1_clip.SD1ClipModel):
         return out, pooled, extra
 
 
-def te(dtype_llama=None, llama_quantization_metadata=None):
+def te(dtype_llama=None, llama_scaled_fp8=None):
     class QwenImageTEModel_(QwenImageTEModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if llama_quantization_metadata is not None:
+            if llama_scaled_fp8 is not None and "scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["quantization_metadata"] = llama_quantization_metadata
+                model_options["scaled_fp8"] = llama_scaled_fp8
             if dtype_llama is not None:
                 dtype = dtype_llama
             super().__init__(device=device, dtype=dtype, model_options=model_options)

comfy/text_encoders/sd3_clip.py

@@ -6,15 +6,14 @@ import torch
 import os
 import comfy.model_management
 import logging
-import comfy.utils
 
 class T5XXLModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, attention_mask=False, model_options={}):
         textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_config_xxl.json")
-        t5xxl_quantization_metadata = model_options.get("t5xxl_quantization_metadata", None)
-        if t5xxl_quantization_metadata is not None:
+        t5xxl_scaled_fp8 = model_options.get("t5xxl_scaled_fp8", None)
+        if t5xxl_scaled_fp8 is not None:
             model_options = model_options.copy()
-            model_options["quantization_metadata"] = t5xxl_quantization_metadata
+            model_options["scaled_fp8"] = t5xxl_scaled_fp8
 
         model_options = {**model_options, "model_name": "t5xxl"}
         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=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
@@ -26,9 +25,9 @@ def t5_xxl_detect(state_dict, prefix=""):
     if t5_key in state_dict:
         out["dtype_t5"] = state_dict[t5_key].dtype
 
-    quant = comfy.utils.detect_layer_quantization(state_dict, prefix)
-    if quant is not None:
-        out["t5_quantization_metadata"] = quant
+    scaled_fp8_key = "{}scaled_fp8".format(prefix)
+    if scaled_fp8_key in state_dict:
+        out["t5xxl_scaled_fp8"] = state_dict[scaled_fp8_key].dtype
 
     return out
 
@@ -157,11 +156,11 @@ class SD3ClipModel(torch.nn.Module):
         else:
             return self.t5xxl.load_sd(sd)
 
-def sd3_clip(clip_l=True, clip_g=True, t5=True, dtype_t5=None, t5_quantization_metadata=None, t5_attention_mask=False):
+def sd3_clip(clip_l=True, clip_g=True, t5=True, dtype_t5=None, t5xxl_scaled_fp8=None, t5_attention_mask=False):
     class SD3ClipModel_(SD3ClipModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if t5_quantization_metadata is not None:
+            if t5xxl_scaled_fp8 is not None and "t5xxl_scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["t5xxl_quantization_metadata"] = t5_quantization_metadata
+                model_options["t5xxl_scaled_fp8"] = t5xxl_scaled_fp8
             super().__init__(clip_l=clip_l, clip_g=clip_g, t5=t5, dtype_t5=dtype_t5, t5_attention_mask=t5_attention_mask, device=device, dtype=dtype, model_options=model_options)
     return SD3ClipModel_

comfy/text_encoders/wan.py

@@ -25,12 +25,12 @@ class WanT5Model(sd1_clip.SD1ClipModel):
     def __init__(self, device="cpu", dtype=None, model_options={}, **kwargs):
         super().__init__(device=device, dtype=dtype, model_options=model_options, name="umt5xxl", clip_model=UMT5XXlModel, **kwargs)
 
-def te(dtype_t5=None, t5_quantization_metadata=None):
+def te(dtype_t5=None, t5xxl_scaled_fp8=None):
     class WanTEModel(WanT5Model):
         def __init__(self, device="cpu", dtype=None, model_options={}):
-            if t5_quantization_metadata is not None:
+            if t5xxl_scaled_fp8 is not None and "scaled_fp8" not in model_options:
                 model_options = model_options.copy()
-                model_options["quantization_metadata"] = t5_quantization_metadata
+                model_options["scaled_fp8"] = t5xxl_scaled_fp8
             if dtype_t5 is not None:
                 dtype = dtype_t5
             super().__init__(device=device, dtype=dtype, model_options=model_options)

comfy/text_encoders/z_image.py

@@ -1,45 +0,0 @@
-from transformers import Qwen2Tokenizer
-import comfy.text_encoders.llama
-from comfy import sd1_clip
-import os
-
-class Qwen3Tokenizer(sd1_clip.SDTokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "qwen25_tokenizer")
-        super().__init__(tokenizer_path, pad_with_end=False, embedding_size=2560, embedding_key='qwen3_4b', tokenizer_class=Qwen2Tokenizer, has_start_token=False, has_end_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, pad_token=151643, tokenizer_data=tokenizer_data)
-
-
-class ZImageTokenizer(sd1_clip.SD1Tokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="qwen3_4b", tokenizer=Qwen3Tokenizer)
-        self.llama_template = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
-
-    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, **kwargs):
-        if llama_template is None:
-            llama_text = self.llama_template.format(text)
-        else:
-            llama_text = llama_template.format(text)
-
-        tokens = super().tokenize_with_weights(llama_text, return_word_ids=return_word_ids, disable_weights=True, **kwargs)
-        return tokens
-
-
-class Qwen3_4BModel(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="hidden", layer_idx=-2, dtype=None, attention_mask=True, model_options={}):
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Qwen3_4B, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
-
-
-class ZImageTEModel(sd1_clip.SD1ClipModel):
-    def __init__(self, device="cpu", dtype=None, model_options={}):
-        super().__init__(device=device, dtype=dtype, name="qwen3_4b", clip_model=Qwen3_4BModel, model_options=model_options)
-
-
-def te(dtype_llama=None, llama_quantization_metadata=None):
-    class ZImageTEModel_(ZImageTEModel):
-        def __init__(self, device="cpu", dtype=None, model_options={}):
-            if dtype_llama is not None:
-                dtype = dtype_llama
-            if llama_quantization_metadata is not None:
-                model_options["quantization_metadata"] = llama_quantization_metadata
-            super().__init__(device=device, dtype=dtype, model_options=model_options)
-    return ZImageTEModel_

comfy/utils.py

@@ -29,7 +29,6 @@ import itertools
 from torch.nn.functional import interpolate
 from einops import rearrange
 from comfy.cli_args import args
-import json
 
 MMAP_TORCH_FILES = args.mmap_torch_files
 DISABLE_MMAP = args.disable_mmap
@@ -53,7 +52,7 @@ if hasattr(torch.serialization, "add_safe_globals"):  # TODO: this was added in
     ALWAYS_SAFE_LOAD = True
     logging.info("Checkpoint files will always be loaded safely.")
 else:
-    logging.warning("Warning, you are using an old pytorch version and some ckpt/pt files might be loaded unsafely. Upgrading to 2.4 or above is recommended as older versions of pytorch are no longer supported.")
+    logging.info("Warning, you are using an old pytorch version and some ckpt/pt files might be loaded unsafely. Upgrading to 2.4 or above is recommended.")
 
 def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False):
     if device is None:
@@ -676,72 +675,6 @@ def flux_to_diffusers(mmdit_config, output_prefix=""):
 
     return key_map
 
-def z_image_to_diffusers(mmdit_config, output_prefix=""):
-    n_layers = mmdit_config.get("n_layers", 0)
-    hidden_size = mmdit_config.get("dim", 0)
-    n_context_refiner = mmdit_config.get("n_refiner_layers", 2)
-    n_noise_refiner = mmdit_config.get("n_refiner_layers", 2)
-    key_map = {}
-
-    def add_block_keys(prefix_from, prefix_to, has_adaln=True):
-        for end in ("weight", "bias"):
-            k = "{}.attention.".format(prefix_from)
-            qkv = "{}.attention.qkv.{}".format(prefix_to, end)
-            key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, hidden_size))
-            key_map["{}to_k.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))
-            key_map["{}to_v.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))
-
-        block_map = {
-            "attention.norm_q.weight": "attention.q_norm.weight",
-            "attention.norm_k.weight": "attention.k_norm.weight",
-            "attention.to_out.0.weight": "attention.out.weight",
-            "attention.to_out.0.bias": "attention.out.bias",
-            "attention_norm1.weight": "attention_norm1.weight",
-            "attention_norm2.weight": "attention_norm2.weight",
-            "feed_forward.w1.weight": "feed_forward.w1.weight",
-            "feed_forward.w2.weight": "feed_forward.w2.weight",
-            "feed_forward.w3.weight": "feed_forward.w3.weight",
-            "ffn_norm1.weight": "ffn_norm1.weight",
-            "ffn_norm2.weight": "ffn_norm2.weight",
-        }
-        if has_adaln:
-            block_map["adaLN_modulation.0.weight"] = "adaLN_modulation.0.weight"
-            block_map["adaLN_modulation.0.bias"] = "adaLN_modulation.0.bias"
-        for k, v in block_map.items():
-            key_map["{}.{}".format(prefix_from, k)] = "{}.{}".format(prefix_to, v)
-
-    for i in range(n_layers):
-        add_block_keys("layers.{}".format(i), "{}layers.{}".format(output_prefix, i))
-
-    for i in range(n_context_refiner):
-        add_block_keys("context_refiner.{}".format(i), "{}context_refiner.{}".format(output_prefix, i))
-
-    for i in range(n_noise_refiner):
-        add_block_keys("noise_refiner.{}".format(i), "{}noise_refiner.{}".format(output_prefix, i))
-
-    MAP_BASIC = [
-        ("final_layer.linear.weight", "all_final_layer.2-1.linear.weight"),
-        ("final_layer.linear.bias", "all_final_layer.2-1.linear.bias"),
-        ("final_layer.adaLN_modulation.1.weight", "all_final_layer.2-1.adaLN_modulation.1.weight"),
-        ("final_layer.adaLN_modulation.1.bias", "all_final_layer.2-1.adaLN_modulation.1.bias"),
-        ("x_embedder.weight", "all_x_embedder.2-1.weight"),
-        ("x_embedder.bias", "all_x_embedder.2-1.bias"),
-        ("x_pad_token", "x_pad_token"),
-        ("cap_embedder.0.weight", "cap_embedder.0.weight"),
-        ("cap_embedder.1.weight", "cap_embedder.1.weight"),
-        ("cap_embedder.1.bias", "cap_embedder.1.bias"),
-        ("cap_pad_token", "cap_pad_token"),
-        ("t_embedder.mlp.0.weight", "t_embedder.mlp.0.weight"),
-        ("t_embedder.mlp.0.bias", "t_embedder.mlp.0.bias"),
-        ("t_embedder.mlp.2.weight", "t_embedder.mlp.2.weight"),
-        ("t_embedder.mlp.2.bias", "t_embedder.mlp.2.bias"),
-    ]
-
-    for c, diffusers in MAP_BASIC:
-        key_map[diffusers] = "{}{}".format(output_prefix, c)
-
-    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)
@@ -803,17 +736,12 @@ def safetensors_header(safetensors_path, max_size=100*1024*1024):
             return None
         return f.read(length_of_header)
 
-ATTR_UNSET={}
-
 def set_attr(obj, attr, value):
     attrs = attr.split(".")
     for name in attrs[:-1]:
         obj = getattr(obj, name)
-    prev = getattr(obj, attrs[-1], ATTR_UNSET)
-    if value is ATTR_UNSET:
-        delattr(obj, attrs[-1])
-    else:
-        setattr(obj, attrs[-1], value)
+    prev = getattr(obj, attrs[-1])
+    setattr(obj, attrs[-1], value)
     return prev
 
 def set_attr_param(obj, attr, value):
@@ -1200,68 +1128,3 @@ def unpack_latents(combined_latent, latent_shapes):
     else:
         output_tensors = combined_latent
     return output_tensors
-
-def detect_layer_quantization(state_dict, prefix):
-    for k in state_dict:
-        if k.startswith(prefix) and k.endswith(".comfy_quant"):
-            logging.info("Found quantization metadata version 1")
-            return {"mixed_ops": True}
-    return None
-
-def convert_old_quants(state_dict, model_prefix="", metadata={}):
-    if metadata is None:
-        metadata = {}
-
-    quant_metadata = None
-    if "_quantization_metadata" not in metadata:
-        scaled_fp8_key = "{}scaled_fp8".format(model_prefix)
-
-        if scaled_fp8_key in state_dict:
-            scaled_fp8_weight = state_dict[scaled_fp8_key]
-            scaled_fp8_dtype = scaled_fp8_weight.dtype
-            if scaled_fp8_dtype == torch.float32:
-                scaled_fp8_dtype = torch.float8_e4m3fn
-
-            if scaled_fp8_weight.nelement() == 2:
-                full_precision_matrix_mult = True
-            else:
-                full_precision_matrix_mult = False
-
-            out_sd = {}
-            layers = {}
-            for k in list(state_dict.keys()):
-                if not k.startswith(model_prefix):
-                    out_sd[k] = state_dict[k]
-                    continue
-                k_out = k
-                w = state_dict.pop(k)
-                layer = None
-                if k_out.endswith(".scale_weight"):
-                    layer = k_out[:-len(".scale_weight")]
-                    k_out = "{}.weight_scale".format(layer)
-
-                if layer is not None:
-                    layer_conf = {"format": "float8_e4m3fn"}  # TODO: check if anyone did some non e4m3fn scaled checkpoints
-                    if full_precision_matrix_mult:
-                        layer_conf["full_precision_matrix_mult"] = full_precision_matrix_mult
-                    layers[layer] = layer_conf
-
-                if k_out.endswith(".scale_input"):
-                    layer = k_out[:-len(".scale_input")]
-                    k_out = "{}.input_scale".format(layer)
-                    if w.item() == 1.0:
-                        continue
-
-                out_sd[k_out] = w
-
-            state_dict = out_sd
-            quant_metadata = {"layers": layers}
-    else:
-        quant_metadata = json.loads(metadata["_quantization_metadata"])
-
-    if quant_metadata is not None:
-        layers = quant_metadata["layers"]
-        for k, v in layers.items():
-            state_dict["{}.comfy_quant".format(k)] = torch.tensor(list(json.dumps(v).encode('utf-8')), dtype=torch.uint8)
-
-    return state_dict, metadata

comfy/weight_adapter/lora.py

@@ -194,7 +194,6 @@ class LoRAAdapter(WeightAdapterBase):
             lora_diff = torch.mm(
                 mat1.flatten(start_dim=1), mat2.flatten(start_dim=1)
             ).reshape(weight.shape)
-            del mat1, mat2
             if dora_scale is not None:
                 weight = weight_decompose(
                     dora_scale,

comfy_api/feature_flags.py

@@ -5,20 +5,19 @@ This module handles capability negotiation between frontend and backend,
 allowing graceful protocol evolution while maintaining backward compatibility.
 """
 
-from typing import Any
+from typing import Any, Dict
 
 from comfy.cli_args import args
 
 # Default server capabilities
-SERVER_FEATURE_FLAGS: dict[str, Any] = {
+SERVER_FEATURE_FLAGS: Dict[str, Any] = {
     "supports_preview_metadata": True,
     "max_upload_size": args.max_upload_size * 1024 * 1024, # Convert MB to bytes
-    "extension": {"manager": {"supports_v4": True}},
 }
 
 
 def get_connection_feature(
-    sockets_metadata: dict[str, dict[str, Any]],
+    sockets_metadata: Dict[str, Dict[str, Any]],
     sid: str,
     feature_name: str,
     default: Any = False
@@ -42,7 +41,7 @@ def get_connection_feature(
 
 
 def supports_feature(
-    sockets_metadata: dict[str, dict[str, Any]],
+    sockets_metadata: Dict[str, Dict[str, Any]],
     sid: str,
     feature_name: str
 ) -> bool:
@@ -60,7 +59,7 @@ def supports_feature(
     return get_connection_feature(sockets_metadata, sid, feature_name, False) is True
 
 
-def get_server_features() -> dict[str, Any]:
+def get_server_features() -> Dict[str, Any]:
     """
     Get the server's feature flags.
 

comfy_api/internal/api_registry.py

@@ -1,4 +1,4 @@
-from typing import NamedTuple
+from typing import Type, List, NamedTuple
 from comfy_api.internal.singleton import ProxiedSingleton
 from packaging import version as packaging_version
 
@@ -10,7 +10,7 @@ class ComfyAPIBase(ProxiedSingleton):
 
 class ComfyAPIWithVersion(NamedTuple):
     version: str
-    api_class: type[ComfyAPIBase]
+    api_class: Type[ComfyAPIBase]
 
 
 def parse_version(version_str: str) -> packaging_version.Version:
@@ -23,16 +23,16 @@ def parse_version(version_str: str) -> packaging_version.Version:
     return packaging_version.parse(version_str)
 
 
-registered_versions: list[ComfyAPIWithVersion] = []
+registered_versions: List[ComfyAPIWithVersion] = []
 
 
-def register_versions(versions: list[ComfyAPIWithVersion]):
+def register_versions(versions: List[ComfyAPIWithVersion]):
     versions.sort(key=lambda x: parse_version(x.version))
     global registered_versions
     registered_versions = versions
 
 
-def get_all_versions() -> list[ComfyAPIWithVersion]:
+def get_all_versions() -> List[ComfyAPIWithVersion]:
     """
     Returns a list of all registered ComfyAPI versions.
     """

comfy_api/internal/async_to_sync.py

@@ -8,7 +8,7 @@ import os
 import textwrap
 import threading
 from enum import Enum
-from typing import Optional, get_origin, get_args, get_type_hints
+from typing import Optional, Type, get_origin, get_args
 
 
 class TypeTracker:
@@ -193,7 +193,7 @@ class AsyncToSyncConverter:
         return result_container["result"]
 
     @classmethod
-    def create_sync_class(cls, async_class: type, thread_pool_size=10) -> type:
+    def create_sync_class(cls, async_class: Type, thread_pool_size=10) -> Type:
         """
         Creates a new class with synchronous versions of all async methods.
 
@@ -220,18 +220,11 @@ class AsyncToSyncConverter:
             self._async_instance = async_class(*args, **kwargs)
 
             # Handle annotated class attributes (like execution: Execution)
-            # Get all annotations from the class hierarchy and resolve string annotations
-            try:
-                # get_type_hints resolves string annotations to actual type objects
-                # This handles classes using 'from __future__ import annotations'
-                all_annotations = get_type_hints(async_class)
-            except Exception:
-                # Fallback to raw annotations if get_type_hints fails
-                # (e.g., for undefined forward references)
-                all_annotations = {}
-                for base_class in reversed(inspect.getmro(async_class)):
-                    if hasattr(base_class, "__annotations__"):
-                        all_annotations.update(base_class.__annotations__)
+            # Get all annotations from the class hierarchy
+            all_annotations = {}
+            for base_class in reversed(inspect.getmro(async_class)):
+                if hasattr(base_class, "__annotations__"):
+                    all_annotations.update(base_class.__annotations__)
 
             # For each annotated attribute, check if it needs to be created or wrapped
             for attr_name, attr_type in all_annotations.items():
@@ -563,7 +556,7 @@ class AsyncToSyncConverter:
 
     @classmethod
     def _generate_imports(
-        cls, async_class: type, type_tracker: TypeTracker
+        cls, async_class: Type, type_tracker: TypeTracker
     ) -> list[str]:
         """Generate import statements for the stub file."""
         imports = []
@@ -628,23 +621,19 @@ class AsyncToSyncConverter:
         return imports
 
     @classmethod
-    def _get_class_attributes(cls, async_class: type) -> list[tuple[str, type]]:
+    def _get_class_attributes(cls, async_class: Type) -> list[tuple[str, Type]]:
         """Extract class attributes that are classes themselves."""
         class_attributes = []
 
-        # Get resolved type hints to handle string annotations
-        try:
-            type_hints = get_type_hints(async_class)
-        except Exception:
-            type_hints = {}
-
         # Look for class attributes that are classes
         for name, attr in sorted(inspect.getmembers(async_class)):
             if isinstance(attr, type) and not name.startswith("_"):
                 class_attributes.append((name, attr))
-            elif name in type_hints:
-                # Use resolved type hint instead of raw annotation
-                annotation = type_hints[name]
+            elif (
+                hasattr(async_class, "__annotations__")
+                and name in async_class.__annotations__
+            ):
+                annotation = async_class.__annotations__[name]
                 if isinstance(annotation, type):
                     class_attributes.append((name, annotation))
 
@@ -654,7 +643,7 @@ class AsyncToSyncConverter:
     def _generate_inner_class_stub(
         cls,
         name: str,
-        attr: type,
+        attr: Type,
         indent: str = "    ",
         type_tracker: Optional[TypeTracker] = None,
     ) -> list[str]:
@@ -782,7 +771,7 @@ class AsyncToSyncConverter:
         return processed
 
     @classmethod
-    def generate_stub_file(cls, async_class: type, sync_class: type) -> None:
+    def generate_stub_file(cls, async_class: Type, sync_class: Type) -> None:
         """
         Generate a .pyi stub file for the sync class to help IDEs with type checking.
         """
@@ -919,15 +908,11 @@ class AsyncToSyncConverter:
             attribute_mappings = {}
 
             # First check annotations for typed attributes (including from parent classes)
-            # Resolve string annotations to actual types
-            try:
-                all_annotations = get_type_hints(async_class)
-            except Exception:
-                # Fallback to raw annotations
-                all_annotations = {}
-                for base_class in reversed(inspect.getmro(async_class)):
-                    if hasattr(base_class, "__annotations__"):
-                        all_annotations.update(base_class.__annotations__)
+            # Collect all annotations from the class hierarchy
+            all_annotations = {}
+            for base_class in reversed(inspect.getmro(async_class)):
+                if hasattr(base_class, "__annotations__"):
+                    all_annotations.update(base_class.__annotations__)
 
             for attr_name, attr_type in sorted(all_annotations.items()):
                 for class_name, class_type in class_attributes:
@@ -988,7 +973,7 @@ class AsyncToSyncConverter:
             logging.error(traceback.format_exc())
 
 
-def create_sync_class(async_class: type, thread_pool_size=10) -> type:
+def create_sync_class(async_class: Type, thread_pool_size=10) -> Type:
     """
     Creates a sync version of an async class
 

comfy_api/internal/singleton.py

@@ -1,4 +1,4 @@
-from typing import TypeVar
+from typing import Type, TypeVar
 
 class SingletonMetaclass(type):
     T = TypeVar("T", bound="SingletonMetaclass")
@@ -11,13 +11,13 @@ class SingletonMetaclass(type):
             )
         return cls._instances[cls]
 
-    def inject_instance(cls: type[T], instance: T) -> None:
+    def inject_instance(cls: Type[T], instance: T) -> None:
         assert cls not in SingletonMetaclass._instances, (
             "Cannot inject instance after first instantiation"
         )
         SingletonMetaclass._instances[cls] = instance
 
-    def get_instance(cls: type[T], *args, **kwargs) -> T:
+    def get_instance(cls: Type[T], *args, **kwargs) -> T:
         """
         Gets the singleton instance of the class, creating it if it doesn't exist.
         """

comfy_api/latest/__init__.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
+from typing import Type, TYPE_CHECKING
 from comfy_api.internal import ComfyAPIBase
 from comfy_api.internal.singleton import ProxiedSingleton
 from comfy_api.internal.async_to_sync import create_sync_class
-from ._input import ImageInput, AudioInput, MaskInput, LatentInput, VideoInput
-from ._input_impl import VideoFromFile, VideoFromComponents
-from ._util import VideoCodec, VideoContainer, VideoComponents, MESH, VOXEL
-from . import _io_public as io
-from . import _ui_public as ui
+from comfy_api.latest._input import ImageInput, AudioInput, MaskInput, LatentInput, VideoInput
+from comfy_api.latest._input_impl import VideoFromFile, VideoFromComponents
+from comfy_api.latest._util import VideoCodec, VideoContainer, VideoComponents
+from . import _io as io
+from . import _ui as ui
 # from comfy_api.latest._resources import _RESOURCES as resources  #noqa: F401
 from comfy_execution.utils import get_executing_context
 from comfy_execution.progress import get_progress_state, PreviewImageTuple
@@ -80,7 +80,7 @@ class ComfyExtension(ABC):
     async def on_load(self) -> None:
         """
         Called when an extension is loaded.
-        This should be used to initialize any global resources needed by the extension.
+        This should be used to initialize any global resources neeeded by the extension.
         """
 
     @abstractmethod
@@ -104,8 +104,6 @@ class Types:
     VideoCodec = VideoCodec
     VideoContainer = VideoContainer
     VideoComponents = VideoComponents
-    MESH = MESH
-    VOXEL = VOXEL
 
 ComfyAPI = ComfyAPI_latest
 
@@ -113,7 +111,7 @@ ComfyAPI = ComfyAPI_latest
 if TYPE_CHECKING:
     import comfy_api.latest.generated.ComfyAPISyncStub  # type: ignore
 
-    ComfyAPISync: type[comfy_api.latest.generated.ComfyAPISyncStub.ComfyAPISyncStub]
+    ComfyAPISync: Type[comfy_api.latest.generated.ComfyAPISyncStub.ComfyAPISyncStub]
 ComfyAPISync = create_sync_class(ComfyAPI_latest)
 
 # create new aliases for io and ui

comfy_api/latest/_input/basic_types.py

@@ -1,5 +1,5 @@
 import torch
-from typing import TypedDict, Optional
+from typing import TypedDict, List, Optional
 
 ImageInput = torch.Tensor
 """
@@ -39,4 +39,4 @@ class LatentInput(TypedDict):
     Optional noise mask tensor in the same format as samples.
     """
 
-    batch_index: Optional[list[int]]
+    batch_index: Optional[List[int]]

comfy_api/latest/_input/video_types.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 from abc import ABC, abstractmethod
-from fractions import Fraction
 from typing import Optional, Union, IO
 import io
 import av
-from .._util import VideoContainer, VideoCodec, VideoComponents
+from comfy_api.util import VideoContainer, VideoCodec, VideoComponents
 
 class VideoInput(ABC):
     """
@@ -73,33 +72,6 @@ class VideoInput(ABC):
         frame_count = components.images.shape[0]
         return float(frame_count / components.frame_rate)
 
-    def get_frame_count(self) -> int:
-        """
-        Returns the number of frames in the video.
-
-        Default implementation uses :meth:`get_components`, which may require
-        loading all frames into memory. File-based implementations should
-        override this method and use container/stream metadata instead.
-
-        Returns:
-            Total number of frames as an integer.
-        """
-        return int(self.get_components().images.shape[0])
-
-    def get_frame_rate(self) -> Fraction:
-        """
-        Returns the frame rate of the video.
-
-        Default implementation materializes the video into memory via
-        `get_components()`. Subclasses that can inspect the underlying
-        container (e.g. `VideoFromFile`) should override this with a more
-        efficient implementation.
-
-        Returns:
-            Frame rate as a Fraction.
-        """
-        return self.get_components().frame_rate
-
     def get_container_format(self) -> str:
         """
         Returns the container format of the video (e.g., 'mp4', 'mov', 'avi').

comfy_api/latest/_input_impl/video_types.py

@@ -3,14 +3,14 @@ from av.container import InputContainer
 from av.subtitles.stream import SubtitleStream
 from fractions import Fraction
 from typing import Optional
-from .._input import AudioInput, VideoInput
+from comfy_api.latest._input import AudioInput, VideoInput
 import av
 import io
 import json
 import numpy as np
 import math
 import torch
-from .._util import VideoContainer, VideoCodec, VideoComponents
+from comfy_api.latest._util import VideoContainer, VideoCodec, VideoComponents
 
 
 def container_to_output_format(container_format: str | None) -> str | None:
@@ -121,71 +121,6 @@ class VideoFromFile(VideoInput):
 
         raise ValueError(f"Could not determine duration for file '{self.__file}'")
 
-    def get_frame_count(self) -> int:
-        """
-        Returns the number of frames in the video without materializing them as
-        torch tensors.
-        """
-        if isinstance(self.__file, io.BytesIO):
-            self.__file.seek(0)
-
-        with av.open(self.__file, mode="r") as container:
-            video_stream = self._get_first_video_stream(container)
-            # 1. Prefer the frames field if available
-            if video_stream.frames and video_stream.frames > 0:
-                return int(video_stream.frames)
-
-            # 2. Try to estimate from duration and average_rate using only metadata
-            if container.duration is not None and video_stream.average_rate:
-                duration_seconds = float(container.duration / av.time_base)
-                estimated_frames = int(round(duration_seconds * float(video_stream.average_rate)))
-                if estimated_frames > 0:
-                    return estimated_frames
-
-            if (
-                getattr(video_stream, "duration", None) is not None
-                and getattr(video_stream, "time_base", None) is not None
-                and video_stream.average_rate
-            ):
-                duration_seconds = float(video_stream.duration * video_stream.time_base)
-                estimated_frames = int(round(duration_seconds * float(video_stream.average_rate)))
-                if estimated_frames > 0:
-                    return estimated_frames
-
-            # 3. Last resort: decode frames and count them (streaming)
-            frame_count = 0
-            container.seek(0)
-            for packet in container.demux(video_stream):
-                for _ in packet.decode():
-                    frame_count += 1
-
-            if frame_count == 0:
-                raise ValueError(f"Could not determine frame count for file '{self.__file}'")
-            return frame_count
-
-    def get_frame_rate(self) -> Fraction:
-        """
-        Returns the average frame rate of the video using container metadata
-        without decoding all frames.
-        """
-        if isinstance(self.__file, io.BytesIO):
-            self.__file.seek(0)
-
-        with av.open(self.__file, mode="r") as container:
-            video_stream = self._get_first_video_stream(container)
-            # Preferred: use PyAV's average_rate (usually already a Fraction-like)
-            if video_stream.average_rate:
-                return Fraction(video_stream.average_rate)
-
-            # Fallback: estimate from frames + duration if available
-            if video_stream.frames and container.duration:
-                duration_seconds = float(container.duration / av.time_base)
-                if duration_seconds > 0:
-                    return Fraction(video_stream.frames / duration_seconds).limit_denominator()
-
-            # Last resort: match get_components_internal default
-            return Fraction(1)
-
     def get_container_format(self) -> str:
         """
         Returns the container format of the video (e.g., 'mp4', 'mov', 'avi').
@@ -303,13 +238,6 @@ class VideoFromFile(VideoInput):
                         packet.stream = stream_map[packet.stream]
                         output_container.mux(packet)
 
-    def _get_first_video_stream(self, container: InputContainer):
-        video_stream = next((s for s in container.streams if s.type == "video"), None)
-        if video_stream is None:
-            raise ValueError(f"No video stream found in file '{self.__file}'")
-        return video_stream
-
-
 class VideoFromComponents(VideoInput):
     """
     Class representing video input from tensors.
@@ -336,10 +264,7 @@ class VideoFromComponents(VideoInput):
             raise ValueError("Only MP4 format is supported for now")
         if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
             raise ValueError("Only H264 codec is supported for now")
-        extra_kwargs = {}
-        if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
-            extra_kwargs["format"] = format.value
-        with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs) as output:
+        with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}) as output:
             # Add metadata before writing any streams
             if metadata is not None:
                 for key, value in metadata.items():

comfy_api/latest/_io.py

@@ -4,8 +4,7 @@ import copy
 import inspect
 from abc import ABC, abstractmethod
 from collections import Counter
-from collections.abc import Iterable
-from dataclasses import asdict, dataclass, field
+from dataclasses import asdict, dataclass
 from enum import Enum
 from typing import Any, Callable, Literal, TypedDict, TypeVar, TYPE_CHECKING
 from typing_extensions import NotRequired, final
@@ -26,9 +25,8 @@ if TYPE_CHECKING:
     from comfy_api.input import VideoInput
 from comfy_api.internal import (_ComfyNodeInternal, _NodeOutputInternal, classproperty, copy_class, first_real_override, is_class,
     prune_dict, shallow_clone_class)
-from ._resources import Resources, ResourcesLocal
+from comfy_api.latest._resources import Resources, ResourcesLocal
 from comfy_execution.graph_utils import ExecutionBlocker
-from ._util import MESH, VOXEL
 
 # from comfy_extras.nodes_images import SVG as SVG_ # NOTE: needs to be moved before can be imported due to circular reference
 
@@ -151,9 +149,6 @@ class _IO_V3:
     def __init__(self):
         pass
 
-    def validate(self):
-        pass
-
     @property
     def io_type(self):
         return self.Parent.io_type
@@ -186,9 +181,6 @@ class Input(_IO_V3):
     def get_io_type(self):
         return _StringIOType(self.io_type)
 
-    def get_all(self) -> list[Input]:
-        return [self]
-
 class WidgetInput(Input):
     '''
     Base class for a V3 Input with widget.
@@ -568,8 +560,6 @@ class Conditioning(ComfyTypeIO):
         '''Used by WAN Camera.'''
         time_dim_concat: NotRequired[torch.Tensor]
         '''Used by WAN Phantom Subject.'''
-        time_dim_replace: NotRequired[torch.Tensor]
-        '''Used by Kandinsky5 I2V.'''
 
     CondList = list[tuple[torch.Tensor, PooledDict]]
     Type = CondList
@@ -638,10 +628,6 @@ class UpscaleModel(ComfyTypeIO):
     if TYPE_CHECKING:
         Type = ImageModelDescriptor
 
-@comfytype(io_type="LATENT_UPSCALE_MODEL")
-class LatentUpscaleModel(ComfyTypeIO):
-    Type = Any
-
 @comfytype(io_type="AUDIO")
 class Audio(ComfyTypeIO):
     class AudioDict(TypedDict):
@@ -670,11 +656,11 @@ class LossMap(ComfyTypeIO):
 
 @comfytype(io_type="VOXEL")
 class Voxel(ComfyTypeIO):
-    Type = VOXEL
+    Type = Any # TODO: VOXEL class is defined in comfy_extras/nodes_hunyuan3d.py; should be moved to somewhere else before referenced directly in v3
 
 @comfytype(io_type="MESH")
 class Mesh(ComfyTypeIO):
-    Type = MESH
+    Type = Any # TODO: MESH class is defined in comfy_extras/nodes_hunyuan3d.py; should be moved to somewhere else before referenced directly in v3
 
 @comfytype(io_type="HOOKS")
 class Hooks(ComfyTypeIO):
@@ -774,13 +760,6 @@ class AudioEncoder(ComfyTypeIO):
 class AudioEncoderOutput(ComfyTypeIO):
     Type = Any
 
-@comfytype(io_type="TRACKS")
-class Tracks(ComfyTypeIO):
-    class TrackDict(TypedDict):
-        track_path: torch.Tensor
-        track_visibility: torch.Tensor
-    Type = TrackDict
-
 @comfytype(io_type="COMFY_MULTITYPED_V3")
 class MultiType:
     Type = Any
@@ -830,61 +809,13 @@ class MultiType:
             else:
                 return super().as_dict()
 
-@comfytype(io_type="COMFY_MATCHTYPE_V3")
-class MatchType(ComfyTypeIO):
-    class Template:
-        def __init__(self, template_id: str, allowed_types: _ComfyType | list[_ComfyType] = AnyType):
-            self.template_id = template_id
-            # account for syntactic sugar
-            if not isinstance(allowed_types, Iterable):
-                allowed_types = [allowed_types]
-            for t in allowed_types:
-                if not isinstance(t, type):
-                    if not isinstance(t, _ComfyType):
-                        raise ValueError(f"Allowed types must be a ComfyType or a list of ComfyTypes, got {t.__class__.__name__}")
-                else:
-                    if not issubclass(t, _ComfyType):
-                        raise ValueError(f"Allowed types must be a ComfyType or a list of ComfyTypes, got {t.__name__}")
-            self.allowed_types = allowed_types
-
-        def as_dict(self):
-            return {
-                "template_id": self.template_id,
-                "allowed_types": ",".join([t.io_type for t in self.allowed_types]),
-            }
-
-    class Input(Input):
-        def __init__(self, id: str, template: MatchType.Template,
-                    display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None):
-            super().__init__(id, display_name, optional, tooltip, lazy, extra_dict)
-            self.template = template
-
-        def as_dict(self):
-            return super().as_dict() | prune_dict({
-                "template": self.template.as_dict(),
-            })
-
-    class Output(Output):
-        def __init__(self, template: MatchType.Template, id: str=None, display_name: str=None, tooltip: str=None,
-                     is_output_list=False):
-            super().__init__(id, display_name, tooltip, is_output_list)
-            self.template = template
-
-        def as_dict(self):
-            return super().as_dict() | prune_dict({
-                "template": self.template.as_dict(),
-            })
-
 class DynamicInput(Input, ABC):
     '''
     Abstract class for dynamic input registration.
     '''
+    @abstractmethod
     def get_dynamic(self) -> list[Input]:
-        return []
-
-    def expand_schema_for_dynamic(self, d: dict[str, Any], live_inputs: dict[str, Any], curr_prefix=''):
-        pass
-
+        ...
 
 class DynamicOutput(Output, ABC):
     '''
@@ -894,223 +825,99 @@ class DynamicOutput(Output, ABC):
                  is_output_list=False):
         super().__init__(id, display_name, tooltip, is_output_list)
 
+    @abstractmethod
     def get_dynamic(self) -> list[Output]:
-        return []
+        ...
 
 
 @comfytype(io_type="COMFY_AUTOGROW_V3")
-class Autogrow(ComfyTypeI):
-    Type = dict[str, Any]
-    _MaxNames = 100  # NOTE: max 100 names for sanity
-
-    class _AutogrowTemplate:
-        def __init__(self, input: Input):
-            # dynamic inputs are not allowed as the template input
-            assert(not isinstance(input, DynamicInput))
-            self.input = copy.copy(input)
-            if isinstance(self.input, WidgetInput):
-                self.input.force_input = True
-            self.names: list[str] = []
-            self.cached_inputs = {}
-
-        def _create_input(self, input: Input, name: str):
-            new_input = copy.copy(self.input)
-            new_input.id = name
-            return new_input
-
-        def _create_cached_inputs(self):
-            for name in self.names:
-                self.cached_inputs[name] = self._create_input(self.input, name)
-
-        def get_all(self) -> list[Input]:
-            return list(self.cached_inputs.values())
-
-        def as_dict(self):
-            return prune_dict({
-                "input": create_input_dict_v1([self.input]),
-            })
-
-        def validate(self):
-            self.input.validate()
-
-        def expand_schema_for_dynamic(self, d: dict[str, Any], live_inputs: dict[str, Any], curr_prefix=''):
-            real_inputs = []
-            for name, input in self.cached_inputs.items():
-                if name in live_inputs:
-                    real_inputs.append(input)
-            add_to_input_dict_v1(d, real_inputs, live_inputs, curr_prefix)
-            add_dynamic_id_mapping(d, real_inputs, curr_prefix)
-
-    class TemplatePrefix(_AutogrowTemplate):
-        def __init__(self, input: Input, prefix: str, min: int=1, max: int=10):
-            super().__init__(input)
-            self.prefix = prefix
-            assert(min >= 0)
-            assert(max >= 1)
-            assert(max <= Autogrow._MaxNames)
-            self.min = min
-            self.max = max
-            self.names = [f"{self.prefix}{i}" for i in range(self.max)]
-            self._create_cached_inputs()
-
-        def as_dict(self):
-            return super().as_dict() | prune_dict({
-                "prefix": self.prefix,
-                "min": self.min,
-                "max": self.max,
-            })
-
-    class TemplateNames(_AutogrowTemplate):
-        def __init__(self, input: Input, names: list[str], min: int=1):
-            super().__init__(input)
-            self.names = names[:Autogrow._MaxNames]
-            assert(min >= 0)
-            self.min = min
-            self._create_cached_inputs()
-
-        def as_dict(self):
-            return super().as_dict() | prune_dict({
-                "names": self.names,
-                "min": self.min,
-            })
-
+class AutogrowDynamic(ComfyTypeI):
+    Type = list[Any]
     class Input(DynamicInput):
-        def __init__(self, id: str, template: Autogrow.TemplatePrefix | Autogrow.TemplateNames,
+        def __init__(self, id: str, template_input: Input, min: int=1, max: int=None,
                      display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None):
             super().__init__(id, display_name, optional, tooltip, lazy, extra_dict)
+            self.template_input = template_input
+            if min is not None:
+                assert(min >= 1)
+            if max is not None:
+                assert(max >= 1)
+            self.min = min
+            self.max = max
+
+        def get_dynamic(self) -> list[Input]:
+            curr_count = 1
+            new_inputs = []
+            for i in range(self.min):
+                new_input = copy.copy(self.template_input)
+                new_input.id = f"{new_input.id}{curr_count}_${self.id}_ag$"
+                if new_input.display_name is not None:
+                    new_input.display_name = f"{new_input.display_name}{curr_count}"
+                new_input.optional = self.optional or new_input.optional
+                if isinstance(self.template_input, WidgetInput):
+                    new_input.force_input = True
+                new_inputs.append(new_input)
+                curr_count += 1
+            # pretend to expand up to max
+            for i in range(curr_count-1, self.max):
+                new_input = copy.copy(self.template_input)
+                new_input.id = f"{new_input.id}{curr_count}_${self.id}_ag$"
+                if new_input.display_name is not None:
+                    new_input.display_name = f"{new_input.display_name}{curr_count}"
+                new_input.optional = True
+                if isinstance(self.template_input, WidgetInput):
+                    new_input.force_input = True
+                new_inputs.append(new_input)
+                curr_count += 1
+            return new_inputs
+
+@comfytype(io_type="COMFY_COMBODYNAMIC_V3")
+class ComboDynamic(ComfyTypeI):
+    class Input(DynamicInput):
+        def __init__(self, id: str):
+            pass
+
+@comfytype(io_type="COMFY_MATCHTYPE_V3")
+class MatchType(ComfyTypeIO):
+    class Template:
+        def __init__(self, template_id: str, allowed_types: _ComfyType | list[_ComfyType]):
+            self.template_id = template_id
+            self.allowed_types = [allowed_types] if isinstance(allowed_types, _ComfyType) else allowed_types
+
+        def as_dict(self):
+            return {
+                "template_id": self.template_id,
+                "allowed_types": "".join(t.io_type for t in self.allowed_types),
+            }
+
+    class Input(DynamicInput):
+        def __init__(self, id: str, template: MatchType.Template,
+                    display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None):
+            super().__init__(id, display_name, optional, tooltip, lazy, extra_dict)
             self.template = template
 
+        def get_dynamic(self) -> list[Input]:
+            return [self]
+
         def as_dict(self):
             return super().as_dict() | prune_dict({
                 "template": self.template.as_dict(),
             })
 
-        def get_dynamic(self) -> list[Input]:
-            return self.template.get_all()
+    class Output(DynamicOutput):
+        def __init__(self, id: str, template: MatchType.Template, display_name: str=None, tooltip: str=None,
+                     is_output_list=False):
+            super().__init__(id, display_name, tooltip, is_output_list)
+            self.template = template
 
-        def get_all(self) -> list[Input]:
-            return [self] + self.template.get_all()
-
-        def validate(self):
-            self.template.validate()
-
-        def expand_schema_for_dynamic(self, d: dict[str, Any], live_inputs: dict[str, Any], curr_prefix=''):
-            curr_prefix = f"{curr_prefix}{self.id}."
-            # need to remove self from expected inputs dictionary; replaced by template inputs in frontend
-            for inner_dict in d.values():
-                if self.id in inner_dict:
-                    del inner_dict[self.id]
-            self.template.expand_schema_for_dynamic(d, live_inputs, curr_prefix)
-
-@comfytype(io_type="COMFY_DYNAMICCOMBO_V3")
-class DynamicCombo(ComfyTypeI):
-    Type = dict[str, Any]
-
-    class Option:
-        def __init__(self, key: str, inputs: list[Input]):
-            self.key = key
-            self.inputs = inputs
-
-        def as_dict(self):
-            return {
-                "key": self.key,
-                "inputs": create_input_dict_v1(self.inputs),
-            }
-
-    class Input(DynamicInput):
-        def __init__(self, id: str, options: list[DynamicCombo.Option],
-                    display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None, extra_dict=None):
-            super().__init__(id, display_name, optional, tooltip, lazy, extra_dict)
-            self.options = options
-
-        def expand_schema_for_dynamic(self, d: dict[str, Any], live_inputs: dict[str, Any], curr_prefix=''):
-            # check if dynamic input's id is in live_inputs
-            if self.id in live_inputs:
-                curr_prefix = f"{curr_prefix}{self.id}."
-                key = live_inputs[self.id]
-                selected_option = None
-                for option in self.options:
-                    if option.key == key:
-                        selected_option = option
-                        break
-                if selected_option is not None:
-                    add_to_input_dict_v1(d, selected_option.inputs, live_inputs, curr_prefix)
-                    add_dynamic_id_mapping(d, selected_option.inputs, curr_prefix, self)
-
-        def get_dynamic(self) -> list[Input]:
-            return [input for option in self.options for input in option.inputs]
-
-        def get_all(self) -> list[Input]:
-            return [self] + [input for option in self.options for input in option.inputs]
+        def get_dynamic(self) -> list[Output]:
+            return [self]
 
         def as_dict(self):
             return super().as_dict() | prune_dict({
-                "options": [o.as_dict() for o in self.options],
+                "template": self.template.as_dict(),
             })
 
-        def validate(self):
-            # make sure all nested inputs are validated
-            for option in self.options:
-                for input in option.inputs:
-                    input.validate()
-
-@comfytype(io_type="COMFY_DYNAMICSLOT_V3")
-class DynamicSlot(ComfyTypeI):
-    Type = dict[str, Any]
-
-    class Input(DynamicInput):
-        def __init__(self, slot: Input, inputs: list[Input],
-                    display_name: str=None, tooltip: str=None, lazy: bool=None, extra_dict=None):
-            assert(not isinstance(slot, DynamicInput))
-            self.slot = copy.copy(slot)
-            self.slot.display_name = slot.display_name if slot.display_name is not None else display_name
-            optional = True
-            self.slot.tooltip = slot.tooltip if slot.tooltip is not None else tooltip
-            self.slot.lazy = slot.lazy if slot.lazy is not None else lazy
-            self.slot.extra_dict = slot.extra_dict if slot.extra_dict is not None else extra_dict
-            super().__init__(slot.id, self.slot.display_name, optional, self.slot.tooltip, self.slot.lazy, self.slot.extra_dict)
-            self.inputs = inputs
-            self.force_input = None
-            # force widget inputs to have no widgets, otherwise this would be awkward
-            if isinstance(self.slot, WidgetInput):
-                self.force_input = True
-                self.slot.force_input = True
-
-        def expand_schema_for_dynamic(self, d: dict[str, Any], live_inputs: dict[str, Any], curr_prefix=''):
-            if self.id in live_inputs:
-                curr_prefix = f"{curr_prefix}{self.id}."
-                add_to_input_dict_v1(d, self.inputs, live_inputs, curr_prefix)
-                add_dynamic_id_mapping(d, [self.slot] + self.inputs, curr_prefix)
-
-        def get_dynamic(self) -> list[Input]:
-            return [self.slot] + self.inputs
-
-        def get_all(self) -> list[Input]:
-            return [self] + [self.slot] + self.inputs
-
-        def as_dict(self):
-            return super().as_dict() | prune_dict({
-                "slotType": str(self.slot.get_io_type()),
-                "inputs": create_input_dict_v1(self.inputs),
-                "forceInput": self.force_input,
-            })
-
-        def validate(self):
-            self.slot.validate()
-            for input in self.inputs:
-                input.validate()
-
-def add_dynamic_id_mapping(d: dict[str, Any], inputs: list[Input], curr_prefix: str, self: DynamicInput=None):
-    dynamic = d.setdefault("dynamic_paths", {})
-    if self is not None:
-        dynamic[self.id] = f"{curr_prefix}{self.id}"
-    for i in inputs:
-        if not isinstance(i, DynamicInput):
-            dynamic[f"{i.id}"] = f"{curr_prefix}{i.id}"
-
-class V3Data(TypedDict):
-    hidden_inputs: dict[str, Any]
-    dynamic_paths: dict[str, Any]
 
 class HiddenHolder:
     def __init__(self, unique_id: str, prompt: Any,
@@ -1172,7 +979,6 @@ class NodeInfoV1:
     output_is_list: list[bool]=None
     output_name: list[str]=None
     output_tooltips: list[str]=None
-    output_matchtypes: list[str]=None
     name: str=None
     display_name: str=None
     description: str=None
@@ -1208,9 +1014,9 @@ class Schema:
     """Display name of node."""
     category: str = "sd"
     """The category of the node, as per the "Add Node" menu."""
-    inputs: list[Input] = field(default_factory=list)
-    outputs: list[Output] = field(default_factory=list)
-    hidden: list[Hidden] = field(default_factory=list)
+    inputs: list[Input]=None
+    outputs: list[Output]=None
+    hidden: list[Hidden]=None
     description: str=""
     """Node description, shown as a tooltip when hovering over the node."""
     is_input_list: bool = False
@@ -1250,11 +1056,7 @@ class Schema:
         '''Validate the schema:
         - verify ids on inputs and outputs are unique - both internally and in relation to each other
         '''
-        nested_inputs: list[Input] = []
-        if self.inputs is not None:
-            for input in self.inputs:
-                nested_inputs.extend(input.get_all())
-        input_ids = [i.id for i in nested_inputs] if nested_inputs is not None else []
+        input_ids = [i.id for i in self.inputs] if self.inputs is not None else []
         output_ids = [o.id for o in self.outputs] if self.outputs is not None else []
         input_set = set(input_ids)
         output_set = set(output_ids)
@@ -1270,13 +1072,6 @@ class Schema:
             issues.append(f"Ids must be unique between inputs and outputs, but {intersection} are not.")
         if len(issues) > 0:
             raise ValueError("\n".join(issues))
-        # validate inputs and outputs
-        if self.inputs is not None:
-            for input in self.inputs:
-                input.validate()
-        if self.outputs is not None:
-            for output in self.outputs:
-                output.validate()
 
     def finalize(self):
         """Add hidden based on selected schema options, and give outputs without ids default ids."""
@@ -1302,10 +1097,19 @@ class Schema:
                 if output.id is None:
                     output.id = f"_{i}_{output.io_type}_"
 
-    def get_v1_info(self, cls, live_inputs: dict[str, Any]=None) -> NodeInfoV1:
-        # NOTE: live_inputs will not be used anymore very soon and this will be done another way
+    def get_v1_info(self, cls) -> NodeInfoV1:
         # get V1 inputs
-        input = create_input_dict_v1(self.inputs, live_inputs)
+        input = {
+            "required": {}
+        }
+        if self.inputs:
+            for i in self.inputs:
+                if isinstance(i, DynamicInput):
+                    dynamic_inputs = i.get_dynamic()
+                    for d in dynamic_inputs:
+                        add_to_dict_v1(d, input)
+                else:
+                    add_to_dict_v1(i, input)
         if self.hidden:
             for hidden in self.hidden:
                 input.setdefault("hidden", {})[hidden.name] = (hidden.value,)
@@ -1314,24 +1118,12 @@ class Schema:
         output_is_list = []
         output_name = []
         output_tooltips = []
-        output_matchtypes = []
-        any_matchtypes = False
         if self.outputs:
             for o in self.outputs:
                 output.append(o.io_type)
                 output_is_list.append(o.is_output_list)
                 output_name.append(o.display_name if o.display_name else o.io_type)
                 output_tooltips.append(o.tooltip if o.tooltip else None)
-                # special handling for MatchType
-                if isinstance(o, MatchType.Output):
-                    output_matchtypes.append(o.template.template_id)
-                    any_matchtypes = True
-                else:
-                    output_matchtypes.append(None)
-
-        # clear out lists that are all None
-        if not any_matchtypes:
-            output_matchtypes = None
 
         info = NodeInfoV1(
             input=input,
@@ -1340,7 +1132,6 @@ class Schema:
             output_is_list=output_is_list,
             output_name=output_name,
             output_tooltips=output_tooltips,
-            output_matchtypes=output_matchtypes,
             name=self.node_id,
             display_name=self.display_name,
             category=self.category,
@@ -1386,57 +1177,16 @@ class Schema:
         return info
 
 
-def create_input_dict_v1(inputs: list[Input], live_inputs: dict[str, Any]=None) -> dict:
-    input = {
-        "required": {}
-    }
-    add_to_input_dict_v1(input, inputs, live_inputs)
-    return input
-
-def add_to_input_dict_v1(d: dict[str, Any], inputs: list[Input], live_inputs: dict[str, Any]=None, curr_prefix=''):
-    for i in inputs:
-        if isinstance(i, DynamicInput):
-            add_to_dict_v1(i, d)
-            if live_inputs is not None:
-                i.expand_schema_for_dynamic(d, live_inputs, curr_prefix)
-        else:
-            add_to_dict_v1(i, d)
-
-def add_to_dict_v1(i: Input, d: dict, dynamic_dict: dict=None):
+def add_to_dict_v1(i: Input, input: dict):
     key = "optional" if i.optional else "required"
     as_dict = i.as_dict()
     # for v1, we don't want to include the optional key
     as_dict.pop("optional", None)
-    if dynamic_dict is None:
-        value = (i.get_io_type(), as_dict)
-    else:
-        value = (i.get_io_type(), as_dict, dynamic_dict)
-    d.setdefault(key, {})[i.id] = value
+    input.setdefault(key, {})[i.id] = (i.get_io_type(), as_dict)
 
 def add_to_dict_v3(io: Input | Output, d: dict):
     d[io.id] = (io.get_io_type(), io.as_dict())
 
-def build_nested_inputs(values: dict[str, Any], v3_data: V3Data):
-    paths = v3_data.get("dynamic_paths", None)
-    if paths is None:
-        return values
-    values = values.copy()
-    result = {}
-
-    for key, path in paths.items():
-        parts = path.split(".")
-        current = result
-
-        for i, p in enumerate(parts):
-            is_last = (i == len(parts) - 1)
-
-            if is_last:
-                current[p] = values.pop(key, None)
-            else:
-                current = current.setdefault(p, {})
-
-    values.update(result)
-    return values
 
 
 class _ComfyNodeBaseInternal(_ComfyNodeInternal):
@@ -1556,12 +1306,12 @@ class _ComfyNodeBaseInternal(_ComfyNodeInternal):
 
     @final
     @classmethod
-    def PREPARE_CLASS_CLONE(cls, v3_data: V3Data | None) -> type[ComfyNode]:
+    def PREPARE_CLASS_CLONE(cls, hidden_inputs: dict) -> type[ComfyNode]:
         """Creates clone of real node class to prevent monkey-patching."""
         c_type: type[ComfyNode] = cls if is_class(cls) else type(cls)
         type_clone: type[ComfyNode] = shallow_clone_class(c_type)
         # set hidden
-        type_clone.hidden = HiddenHolder.from_dict(v3_data["hidden_inputs"] if v3_data else None)
+        type_clone.hidden = HiddenHolder.from_dict(hidden_inputs)
         return type_clone
 
     @final
@@ -1678,18 +1428,14 @@ class _ComfyNodeBaseInternal(_ComfyNodeInternal):
 
     @final
     @classmethod
-    def INPUT_TYPES(cls, include_hidden=True, return_schema=False, live_inputs=None) -> dict[str, dict] | tuple[dict[str, dict], Schema, V3Data]:
+    def INPUT_TYPES(cls, include_hidden=True, return_schema=False) -> dict[str, dict] | tuple[dict[str, dict], Schema]:
         schema = cls.FINALIZE_SCHEMA()
-        info = schema.get_v1_info(cls, live_inputs)
+        info = schema.get_v1_info(cls)
         input = info.input
         if not include_hidden:
             input.pop("hidden", None)
         if return_schema:
-            v3_data: V3Data = {}
-            dynamic = input.pop("dynamic_paths", None)
-            if dynamic is not None:
-                v3_data["dynamic_paths"] = dynamic
-            return input, schema, v3_data
+            return input, schema
         return input
 
     @final
@@ -1762,7 +1508,7 @@ class ComfyNode(_ComfyNodeBaseInternal):
         raise NotImplementedError
 
     @classmethod
-    def validate_inputs(cls, **kwargs) -> bool | str:
+    def validate_inputs(cls, **kwargs) -> bool:
         """Optionally, define this function to validate inputs; equivalent to V1's VALIDATE_INPUTS."""
         raise NotImplementedError
 
@@ -1822,7 +1568,7 @@ class NodeOutput(_NodeOutputInternal):
             ui = data["ui"]
         if "expand" in data:
             expand = data["expand"]
-        return cls(*args, ui=ui, expand=expand)
+        return cls(args=args, ui=ui, expand=expand)
 
     def __getitem__(self, index) -> Any:
         return self.args[index]
@@ -1877,7 +1623,6 @@ __all__ = [
     "StyleModel",
     "Gligen",
     "UpscaleModel",
-    "LatentUpscaleModel",
     "Audio",
     "Video",
     "SVG",
@@ -1901,11 +1646,6 @@ __all__ = [
     "SEGS",
     "AnyType",
     "MultiType",
-    "Tracks",
-    # Dynamic Types
-    "MatchType",
-    # "DynamicCombo",
-    # "Autogrow",
     # Other classes
     "HiddenHolder",
     "Hidden",
@@ -1916,5 +1656,4 @@ __all__ = [
     "NodeOutput",
     "add_to_dict_v1",
     "add_to_dict_v3",
-    "V3Data",
 ]

comfy_api/latest/_io_public.py

@@ -1 +0,0 @@
-from ._io import *  # noqa: F403

comfy_api/latest/_ui.py

@@ -3,8 +3,8 @@ from __future__ import annotations
 import json
 import os
 import random
-import uuid
 from io import BytesIO
+from typing import Type
 
 import av
 import numpy as np
@@ -21,7 +21,7 @@ import folder_paths
 
 # used for image preview
 from comfy.cli_args import args
-from ._io import ComfyNode, FolderType, Image, _UIOutput
+from comfy_api.latest._io import ComfyNode, FolderType, Image, _UIOutput
 
 
 class SavedResult(dict):
@@ -82,7 +82,7 @@ class ImageSaveHelper:
         return PILImage.fromarray(np.clip(255.0 * image_tensor.cpu().numpy(), 0, 255).astype(np.uint8))
 
     @staticmethod
-    def _create_png_metadata(cls: type[ComfyNode] | None) -> PngInfo | None:
+    def _create_png_metadata(cls: Type[ComfyNode] | None) -> PngInfo | None:
         """Creates a PngInfo object with prompt and extra_pnginfo."""
         if args.disable_metadata or cls is None or not cls.hidden:
             return None
@@ -95,7 +95,7 @@ class ImageSaveHelper:
         return metadata
 
     @staticmethod
-    def _create_animated_png_metadata(cls: type[ComfyNode] | None) -> PngInfo | None:
+    def _create_animated_png_metadata(cls: Type[ComfyNode] | None) -> PngInfo | None:
         """Creates a PngInfo object with prompt and extra_pnginfo for animated PNGs (APNG)."""
         if args.disable_metadata or cls is None or not cls.hidden:
             return None
@@ -120,7 +120,7 @@ class ImageSaveHelper:
         return metadata
 
     @staticmethod
-    def _create_webp_metadata(pil_image: PILImage.Image, cls: type[ComfyNode] | None) -> PILImage.Exif:
+    def _create_webp_metadata(pil_image: PILImage.Image, cls: Type[ComfyNode] | None) -> PILImage.Exif:
         """Creates EXIF metadata bytes for WebP images."""
         exif_data = pil_image.getexif()
         if args.disable_metadata or cls is None or cls.hidden is None:
@@ -136,7 +136,7 @@ class ImageSaveHelper:
 
     @staticmethod
     def save_images(
-        images, filename_prefix: str, folder_type: FolderType, cls: type[ComfyNode] | None, compress_level = 4,
+        images, filename_prefix: str, folder_type: FolderType, cls: Type[ComfyNode] | None, compress_level = 4,
     ) -> list[SavedResult]:
         """Saves a batch of images as individual PNG files."""
         full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
@@ -154,7 +154,7 @@ class ImageSaveHelper:
         return results
 
     @staticmethod
-    def get_save_images_ui(images, filename_prefix: str, cls: type[ComfyNode] | None, compress_level=4) -> SavedImages:
+    def get_save_images_ui(images, filename_prefix: str, cls: Type[ComfyNode] | None, compress_level=4) -> SavedImages:
         """Saves a batch of images and returns a UI object for the node output."""
         return SavedImages(
                 ImageSaveHelper.save_images(
@@ -168,7 +168,7 @@ class ImageSaveHelper:
 
     @staticmethod
     def save_animated_png(
-        images, filename_prefix: str, folder_type: FolderType, cls: type[ComfyNode] | None, fps: float, compress_level: int
+        images, filename_prefix: str, folder_type: FolderType, cls: Type[ComfyNode] | None, fps: float, compress_level: int
     ) -> SavedResult:
         """Saves a batch of images as a single animated PNG."""
         full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
@@ -190,7 +190,7 @@ class ImageSaveHelper:
 
     @staticmethod
     def get_save_animated_png_ui(
-        images, filename_prefix: str, cls: type[ComfyNode] | None, fps: float, compress_level: int
+        images, filename_prefix: str, cls: Type[ComfyNode] | None, fps: float, compress_level: int
     ) -> SavedImages:
         """Saves an animated PNG and returns a UI object for the node output."""
         result = ImageSaveHelper.save_animated_png(
@@ -208,7 +208,7 @@ class ImageSaveHelper:
         images,
         filename_prefix: str,
         folder_type: FolderType,
-        cls: type[ComfyNode] | None,
+        cls: Type[ComfyNode] | None,
         fps: float,
         lossless: bool,
         quality: int,
@@ -237,7 +237,7 @@ class ImageSaveHelper:
     def get_save_animated_webp_ui(
         images,
         filename_prefix: str,
-        cls: type[ComfyNode] | None,
+        cls: Type[ComfyNode] | None,
         fps: float,
         lossless: bool,
         quality: int,
@@ -266,7 +266,7 @@ class AudioSaveHelper:
         audio: dict,
         filename_prefix: str,
         folder_type: FolderType,
-        cls: type[ComfyNode] | None,
+        cls: Type[ComfyNode] | None,
         format: str = "flac",
         quality: str = "128k",
     ) -> list[SavedResult]:
@@ -318,10 +318,9 @@ class AudioSaveHelper:
             for key, value in metadata.items():
                 output_container.metadata[key] = value
 
-            layout = "mono" if waveform.shape[0] == 1 else "stereo"
             # Set up the output stream with appropriate properties
             if format == "opus":
-                out_stream = output_container.add_stream("libopus", rate=sample_rate, layout=layout)
+                out_stream = output_container.add_stream("libopus", rate=sample_rate)
                 if quality == "64k":
                     out_stream.bit_rate = 64000
                 elif quality == "96k":
@@ -333,7 +332,7 @@ class AudioSaveHelper:
                 elif quality == "320k":
                     out_stream.bit_rate = 320000
             elif format == "mp3":
-                out_stream = output_container.add_stream("libmp3lame", rate=sample_rate, layout=layout)
+                out_stream = output_container.add_stream("libmp3lame", rate=sample_rate)
                 if quality == "V0":
                     # TODO i would really love to support V3 and V5 but there doesn't seem to be a way to set the qscale level, the property below is a bool
                     out_stream.codec_context.qscale = 1
@@ -342,12 +341,12 @@ class AudioSaveHelper:
                 elif quality == "320k":
                     out_stream.bit_rate = 320000
             else:  # format == "flac":
-                out_stream = output_container.add_stream("flac", rate=sample_rate, layout=layout)
+                out_stream = output_container.add_stream("flac", rate=sample_rate)
 
             frame = av.AudioFrame.from_ndarray(
                 waveform.movedim(0, 1).reshape(1, -1).float().numpy(),
                 format="flt",
-                layout=layout,
+                layout="mono" if waveform.shape[0] == 1 else "stereo",
             )
             frame.sample_rate = sample_rate
             frame.pts = 0
@@ -371,7 +370,7 @@ class AudioSaveHelper:
 
     @staticmethod
     def get_save_audio_ui(
-        audio, filename_prefix: str, cls: type[ComfyNode] | None, format: str = "flac", quality: str = "128k",
+        audio, filename_prefix: str, cls: Type[ComfyNode] | None, format: str = "flac", quality: str = "128k",
     ) -> SavedAudios:
         """Save and instantly wrap for UI."""
         return SavedAudios(
@@ -387,7 +386,7 @@ class AudioSaveHelper:
 
 
 class PreviewImage(_UIOutput):
-    def __init__(self, image: Image.Type, animated: bool = False, cls: type[ComfyNode] = None, **kwargs):
+    def __init__(self, image: Image.Type, animated: bool = False, cls: Type[ComfyNode] = None, **kwargs):
         self.values = ImageSaveHelper.save_images(
             image,
             filename_prefix="ComfyUI_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for _ in range(5)),
@@ -411,7 +410,7 @@ class PreviewMask(PreviewImage):
 
 
 class PreviewAudio(_UIOutput):
-    def __init__(self, audio: dict, cls: type[ComfyNode] = None, **kwargs):
+    def __init__(self, audio: dict, cls: Type[ComfyNode] = None, **kwargs):
         self.values = AudioSaveHelper.save_audio(
             audio,
             filename_prefix="ComfyUI_temp_" + "".join(random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(5)),
@@ -437,19 +436,9 @@ class PreviewUI3D(_UIOutput):
     def __init__(self, model_file, camera_info, **kwargs):
         self.model_file = model_file
         self.camera_info = camera_info
-        self.bg_image_path = None
-        bg_image = kwargs.get("bg_image", None)
-        if bg_image is not None:
-            img_array = (bg_image[0].cpu().numpy() * 255).astype(np.uint8)
-            img = PILImage.fromarray(img_array)
-            temp_dir = folder_paths.get_temp_directory()
-            filename = f"bg_{uuid.uuid4().hex}.png"
-            bg_image_path = os.path.join(temp_dir, filename)
-            img.save(bg_image_path, compress_level=1)
-            self.bg_image_path = f"temp/{filename}"
 
     def as_dict(self):
-        return {"result": [self.model_file, self.camera_info, self.bg_image_path]}
+        return {"result": [self.model_file, self.camera_info]}
 
 
 class PreviewText(_UIOutput):

comfy_api/latest/_ui_public.py

@@ -1 +0,0 @@
-from ._ui import *  # noqa: F403

comfy_api/latest/_util/__init__.py

@@ -1,11 +1,8 @@
 from .video_types import VideoContainer, VideoCodec, VideoComponents
-from .geometry_types import VOXEL, MESH
 
 __all__ = [
     # Utility Types
     "VideoContainer",
     "VideoCodec",
     "VideoComponents",
-    "VOXEL",
-    "MESH",
 ]

comfy_api/latest/_util/geometry_types.py

@@ -1,12 +0,0 @@
-import torch
-
-
-class VOXEL:
-    def __init__(self, data: torch.Tensor):
-        self.data = data
-
-
-class MESH:
-    def __init__(self, vertices: torch.Tensor, faces: torch.Tensor):
-        self.vertices = vertices
-        self.faces = faces

comfy_api/latest/_util/video_types.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from enum import Enum
 from fractions import Fraction
 from typing import Optional
-from .._input import ImageInput, AudioInput
+from comfy_api.latest._input import ImageInput, AudioInput
 
 class VideoCodec(str, Enum):
     AUTO = "auto"

comfy_api/v0_0_2/__init__.py

@@ -6,7 +6,7 @@ from comfy_api.latest import (
 )
 from typing import Type, TYPE_CHECKING
 from comfy_api.internal.async_to_sync import create_sync_class
-from comfy_api.latest import io, ui, IO, UI, ComfyExtension  #noqa: F401
+from comfy_api.latest import io, ui, ComfyExtension  #noqa: F401
 
 
 class ComfyAPIAdapter_v0_0_2(ComfyAPI_latest):
@@ -42,8 +42,4 @@ __all__ = [
     "InputImpl",
     "Types",
     "ComfyExtension",
-    "io",
-    "IO",
-    "ui",
-    "UI",
 ]

comfy_api/version_list.py

@@ -2,8 +2,9 @@ from comfy_api.latest import ComfyAPI_latest
 from comfy_api.v0_0_2 import ComfyAPIAdapter_v0_0_2
 from comfy_api.v0_0_1 import ComfyAPIAdapter_v0_0_1
 from comfy_api.internal import ComfyAPIBase
+from typing import List, Type
 
-supported_versions: list[type[ComfyAPIBase]] = [
+supported_versions: List[Type[ComfyAPIBase]] = [
     ComfyAPI_latest,
     ComfyAPIAdapter_v0_0_2,
     ComfyAPIAdapter_v0_0_1,

comfy_api_nodes/apis/bfl_api.py

@@ -70,29 +70,6 @@ class BFLFluxProGenerateRequest(BaseModel):
     # )
 
 
-class Flux2ProGenerateRequest(BaseModel):
-    prompt: str = Field(...)
-    width: int = Field(1024, description="Must be a multiple of 32.")
-    height: int = Field(768, description="Must be a multiple of 32.")
-    seed: int | None = Field(None)
-    prompt_upsampling: bool | None = Field(None)
-    input_image: str | None = Field(None, description="Base64 encoded image for image-to-image generation")
-    input_image_2: str | None = Field(None, description="Base64 encoded image for image-to-image generation")
-    input_image_3: str | None = Field(None, description="Base64 encoded image for image-to-image generation")
-    input_image_4: str | None = Field(None, description="Base64 encoded image for image-to-image generation")
-    input_image_5: str | None = Field(None, description="Base64 encoded image for image-to-image generation")
-    input_image_6: str | None = Field(None, description="Base64 encoded image for image-to-image generation")
-    input_image_7: str | None = Field(None, description="Base64 encoded image for image-to-image generation")
-    input_image_8: str | None = Field(None, description="Base64 encoded image for image-to-image generation")
-    input_image_9: str | None = Field(None, description="Base64 encoded image for image-to-image generation")
-    safety_tolerance: int | None = Field(
-        5, description="Tolerance level for input and output moderation. Value 0 being most strict.", ge=0, le=5
-    )
-    output_format: str | None = Field(
-        "png", description="Output format for the generated image. Can be 'jpeg' or 'png'."
-    )
-
-
 class BFLFluxKontextProGenerateRequest(BaseModel):
     prompt: str = Field(..., description='The text prompt for what you wannt to edit.')
     input_image: Optional[str] = Field(None, description='Image to edit in base64 format')
@@ -132,9 +109,8 @@ class BFLFluxProUltraGenerateRequest(BaseModel):
 
 
 class BFLFluxProGenerateResponse(BaseModel):
-    id: str = Field(..., description="The unique identifier for the generation task.")
-    polling_url: str = Field(..., description="URL to poll for the generation result.")
-    cost: float | None = Field(None, description="Price in cents")
+    id: str = Field(..., description='The unique identifier for the generation task.')
+    polling_url: str = Field(..., description='URL to poll for the generation result.')
 
 
 class BFLStatus(str, Enum):

comfy_api_nodes/apis/bytedance_api.py

@@ -1,144 +0,0 @@
-from typing import Literal
-
-from pydantic import BaseModel, Field
-
-
-class Text2ImageTaskCreationRequest(BaseModel):
-    model: str = Field(...)
-    prompt: str = Field(...)
-    response_format: str | None = Field("url")
-    size: str | None = Field(None)
-    seed: int | None = Field(0, ge=0, le=2147483647)
-    guidance_scale: float | None = Field(..., ge=1.0, le=10.0)
-    watermark: bool | None = Field(False)
-
-
-class Image2ImageTaskCreationRequest(BaseModel):
-    model: str = Field(...)
-    prompt: str = Field(...)
-    response_format: str | None = Field("url")
-    image: str = Field(..., description="Base64 encoded string or image URL")
-    size: str | None = Field("adaptive")
-    seed: int | None = Field(..., ge=0, le=2147483647)
-    guidance_scale: float | None = Field(..., ge=1.0, le=10.0)
-    watermark: bool | None = Field(False)
-
-
-class Seedream4Options(BaseModel):
-    max_images: int = Field(15)
-
-
-class Seedream4TaskCreationRequest(BaseModel):
-    model: str = Field(...)
-    prompt: str = Field(...)
-    response_format: str = Field("url")
-    image: list[str] | None = Field(None, description="Image URLs")
-    size: str = Field(...)
-    seed: int = Field(..., ge=0, le=2147483647)
-    sequential_image_generation: str = Field("disabled")
-    sequential_image_generation_options: Seedream4Options = Field(Seedream4Options(max_images=15))
-    watermark: bool = Field(False)
-
-
-class ImageTaskCreationResponse(BaseModel):
-    model: str = Field(...)
-    created: int = Field(..., description="Unix timestamp (in seconds) indicating time when the request was created.")
-    data: list = Field([], description="Contains information about the generated image(s).")
-    error: dict = Field({}, description="Contains `code` and `message` fields in case of error.")
-
-
-class TaskTextContent(BaseModel):
-    type: str = Field("text")
-    text: str = Field(...)
-
-
-class TaskImageContentUrl(BaseModel):
-    url: str = Field(...)
-
-
-class TaskImageContent(BaseModel):
-    type: str = Field("image_url")
-    image_url: TaskImageContentUrl = Field(...)
-    role: Literal["first_frame", "last_frame", "reference_image"] | None = Field(None)
-
-
-class Text2VideoTaskCreationRequest(BaseModel):
-    model: str = Field(...)
-    content: list[TaskTextContent] = Field(..., min_length=1)
-
-
-class Image2VideoTaskCreationRequest(BaseModel):
-    model: str = Field(...)
-    content: list[TaskTextContent | TaskImageContent] = Field(..., min_length=2)
-
-
-class TaskCreationResponse(BaseModel):
-    id: str = Field(...)
-
-
-class TaskStatusError(BaseModel):
-    code: str = Field(...)
-    message: str = Field(...)
-
-
-class TaskStatusResult(BaseModel):
-    video_url: str = Field(...)
-
-
-class TaskStatusResponse(BaseModel):
-    id: str = Field(...)
-    model: str = Field(...)
-    status: Literal["queued", "running", "cancelled", "succeeded", "failed"] = Field(...)
-    error: TaskStatusError | None = Field(None)
-    content: TaskStatusResult | None = Field(None)
-
-
-RECOMMENDED_PRESETS = [
-    ("1024x1024 (1:1)", 1024, 1024),
-    ("864x1152 (3:4)", 864, 1152),
-    ("1152x864 (4:3)", 1152, 864),
-    ("1280x720 (16:9)", 1280, 720),
-    ("720x1280 (9:16)", 720, 1280),
-    ("832x1248 (2:3)", 832, 1248),
-    ("1248x832 (3:2)", 1248, 832),
-    ("1512x648 (21:9)", 1512, 648),
-    ("2048x2048 (1:1)", 2048, 2048),
-    ("Custom", None, None),
-]
-
-RECOMMENDED_PRESETS_SEEDREAM_4 = [
-    ("2048x2048 (1:1)", 2048, 2048),
-    ("2304x1728 (4:3)", 2304, 1728),
-    ("1728x2304 (3:4)", 1728, 2304),
-    ("2560x1440 (16:9)", 2560, 1440),
-    ("1440x2560 (9:16)", 1440, 2560),
-    ("2496x1664 (3:2)", 2496, 1664),
-    ("1664x2496 (2:3)", 1664, 2496),
-    ("3024x1296 (21:9)", 3024, 1296),
-    ("4096x4096 (1:1)", 4096, 4096),
-    ("Custom", None, None),
-]
-
-# The time in this dictionary are given for 10 seconds duration.
-VIDEO_TASKS_EXECUTION_TIME = {
-    "seedance-1-0-lite-t2v-250428": {
-        "480p": 40,
-        "720p": 60,
-        "1080p": 90,
-    },
-    "seedance-1-0-lite-i2v-250428": {
-        "480p": 40,
-        "720p": 60,
-        "1080p": 90,
-    },
-    "seedance-1-0-pro-250528": {
-        "480p": 70,
-        "720p": 85,
-        "1080p": 115,
-    },
-    "seedance-1-0-pro-fast-251015": {
-        "480p": 50,
-        "720p": 65,
-        "1080p": 100,
-    },
-}

comfy_api_nodes/apis/gemini_api.py

@@ -1,229 +1,22 @@
-from datetime import date
-from enum import Enum
-from typing import Any
+from typing import Optional
 
-from pydantic import BaseModel, Field
-
-
-class GeminiSafetyCategory(str, Enum):
-    HARM_CATEGORY_SEXUALLY_EXPLICIT = "HARM_CATEGORY_SEXUALLY_EXPLICIT"
-    HARM_CATEGORY_HATE_SPEECH = "HARM_CATEGORY_HATE_SPEECH"
-    HARM_CATEGORY_HARASSMENT = "HARM_CATEGORY_HARASSMENT"
-    HARM_CATEGORY_DANGEROUS_CONTENT = "HARM_CATEGORY_DANGEROUS_CONTENT"
-
-
-class GeminiSafetyThreshold(str, Enum):
-    OFF = "OFF"
-    BLOCK_NONE = "BLOCK_NONE"
-    BLOCK_LOW_AND_ABOVE = "BLOCK_LOW_AND_ABOVE"
-    BLOCK_MEDIUM_AND_ABOVE = "BLOCK_MEDIUM_AND_ABOVE"
-    BLOCK_ONLY_HIGH = "BLOCK_ONLY_HIGH"
-
-
-class GeminiSafetySetting(BaseModel):
-    category: GeminiSafetyCategory
-    threshold: GeminiSafetyThreshold
-
-
-class GeminiRole(str, Enum):
-    user = "user"
-    model = "model"
-
-
-class GeminiMimeType(str, Enum):
-    application_pdf = "application/pdf"
-    audio_mpeg = "audio/mpeg"
-    audio_mp3 = "audio/mp3"
-    audio_wav = "audio/wav"
-    image_png = "image/png"
-    image_jpeg = "image/jpeg"
-    image_webp = "image/webp"
-    text_plain = "text/plain"
-    video_mov = "video/mov"
-    video_mpeg = "video/mpeg"
-    video_mp4 = "video/mp4"
-    video_mpg = "video/mpg"
-    video_avi = "video/avi"
-    video_wmv = "video/wmv"
-    video_mpegps = "video/mpegps"
-    video_flv = "video/flv"
-
-
-class GeminiInlineData(BaseModel):
-    data: str | None = Field(
-        None,
-        description="The base64 encoding of the image, PDF, or video to include inline in the prompt. "
-        "When including media inline, you must also specify the media type (mimeType) of the data. Size limit: 20MB",
-    )
-    mimeType: GeminiMimeType | None = Field(None)
-
-
-class GeminiFileData(BaseModel):
-    fileUri: str | None = Field(None)
-    mimeType: GeminiMimeType | None = Field(None)
-
-
-class GeminiPart(BaseModel):
-    inlineData: GeminiInlineData | None = Field(None)
-    fileData: GeminiFileData | None = Field(None)
-    text: str | None = Field(None)
-
-
-class GeminiTextPart(BaseModel):
-    text: str | None = Field(None)
-
-
-class GeminiContent(BaseModel):
-    parts: list[GeminiPart] = Field([])
-    role: GeminiRole = Field(..., examples=["user"])
-
-
-class GeminiSystemInstructionContent(BaseModel):
-    parts: list[GeminiTextPart] = Field(
-        ...,
-        description="A list of ordered parts that make up a single message. "
-        "Different parts may have different IANA MIME types.",
-    )
-    role: GeminiRole | None = Field(..., description="The role field of systemInstruction may be ignored.")
-
-
-class GeminiFunctionDeclaration(BaseModel):
-    description: str | None = Field(None)
-    name: str = Field(...)
-    parameters: dict[str, Any] = Field(..., description="JSON schema for the function parameters")
-
-
-class GeminiTool(BaseModel):
-    functionDeclarations: list[GeminiFunctionDeclaration] | None = Field(None)
-
-
-class GeminiOffset(BaseModel):
-    nanos: int | None = Field(None, ge=0, le=999999999)
-    seconds: int | None = Field(None, ge=-315576000000, le=315576000000)
-
-
-class GeminiVideoMetadata(BaseModel):
-    endOffset: GeminiOffset | None = Field(None)
-    startOffset: GeminiOffset | None = Field(None)
-
-
-class GeminiGenerationConfig(BaseModel):
-    maxOutputTokens: int | None = Field(None, ge=16, le=8192)
-    seed: int | None = Field(None)
-    stopSequences: list[str] | None = Field(None)
-    temperature: float | None = Field(None, ge=0.0, le=2.0)
-    topK: int | None = Field(None, ge=1)
-    topP: float | None = Field(None, ge=0.0, le=1.0)
+from comfy_api_nodes.apis import GeminiGenerationConfig, GeminiContent, GeminiSafetySetting, GeminiSystemInstructionContent, GeminiTool, GeminiVideoMetadata
+from pydantic import BaseModel
 
 
 class GeminiImageConfig(BaseModel):
-    aspectRatio: str | None = Field(None)
-    imageSize: str | None = Field(None)
+    aspectRatio: Optional[str] = None
 
 
 class GeminiImageGenerationConfig(GeminiGenerationConfig):
-    responseModalities: list[str] | None = Field(None)
-    imageConfig: GeminiImageConfig | None = Field(None)
+    responseModalities: Optional[list[str]] = None
+    imageConfig: Optional[GeminiImageConfig] = None
 
 
 class GeminiImageGenerateContentRequest(BaseModel):
-    contents: list[GeminiContent] = Field(...)
-    generationConfig: GeminiImageGenerationConfig | None = Field(None)
-    safetySettings: list[GeminiSafetySetting] | None = Field(None)
-    systemInstruction: GeminiSystemInstructionContent | None = Field(None)
-    tools: list[GeminiTool] | None = Field(None)
-    videoMetadata: GeminiVideoMetadata | None = Field(None)
-    uploadImagesToStorage: bool = Field(True)
-
-
-class GeminiGenerateContentRequest(BaseModel):
-    contents: list[GeminiContent] = Field(...)
-    generationConfig: GeminiGenerationConfig | None = Field(None)
-    safetySettings: list[GeminiSafetySetting] | None = Field(None)
-    systemInstruction: GeminiSystemInstructionContent | None = Field(None)
-    tools: list[GeminiTool] | None = Field(None)
-    videoMetadata: GeminiVideoMetadata | None = Field(None)
-
-
-class Modality(str, Enum):
-    MODALITY_UNSPECIFIED = "MODALITY_UNSPECIFIED"
-    TEXT = "TEXT"
-    IMAGE = "IMAGE"
-    VIDEO = "VIDEO"
-    AUDIO = "AUDIO"
-    DOCUMENT = "DOCUMENT"
-
-
-class ModalityTokenCount(BaseModel):
-    modality: Modality | None = None
-    tokenCount: int | None = Field(None, description="Number of tokens for the given modality.")
-
-
-class Probability(str, Enum):
-    NEGLIGIBLE = "NEGLIGIBLE"
-    LOW = "LOW"
-    MEDIUM = "MEDIUM"
-    HIGH = "HIGH"
-    UNKNOWN = "UNKNOWN"
-
-
-class GeminiSafetyRating(BaseModel):
-    category: GeminiSafetyCategory | None = None
-    probability: Probability | None = Field(
-        None,
-        description="The probability that the content violates the specified safety category",
-    )
-
-
-class GeminiCitation(BaseModel):
-    authors: list[str] | None = None
-    endIndex: int | None = None
-    license: str | None = None
-    publicationDate: date | None = None
-    startIndex: int | None = None
-    title: str | None = None
-    uri: str | None = None
-
-
-class GeminiCitationMetadata(BaseModel):
-    citations: list[GeminiCitation] | None = None
-
-
-class GeminiCandidate(BaseModel):
-    citationMetadata: GeminiCitationMetadata | None = None
-    content: GeminiContent | None = None
-    finishReason: str | None = None
-    safetyRatings: list[GeminiSafetyRating] | None = None
-
-
-class GeminiPromptFeedback(BaseModel):
-    blockReason: str | None = None
-    blockReasonMessage: str | None = None
-    safetyRatings: list[GeminiSafetyRating] | None = None
-
-
-class GeminiUsageMetadata(BaseModel):
-    cachedContentTokenCount: int | None = Field(
-        None,
-        description="Output only. Number of tokens in the cached part in the input (the cached content).",
-    )
-    candidatesTokenCount: int | None = Field(None, description="Number of tokens in the response(s).")
-    candidatesTokensDetails: list[ModalityTokenCount] | None = Field(
-        None, description="Breakdown of candidate tokens by modality."
-    )
-    promptTokenCount: int | None = Field(
-        None,
-        description="Number of tokens in the request. When cachedContent is set, this is still the total effective prompt size meaning this includes the number of tokens in the cached content.",
-    )
-    promptTokensDetails: list[ModalityTokenCount] | None = Field(
-        None, description="Breakdown of prompt tokens by modality."
-    )
-    thoughtsTokenCount: int | None = Field(None, description="Number of tokens present in thoughts output.")
-    toolUsePromptTokenCount: int | None = Field(None, description="Number of tokens present in tool-use prompt(s).")
-
-
-class GeminiGenerateContentResponse(BaseModel):
-    candidates: list[GeminiCandidate] | None = Field(None)
-    promptFeedback: GeminiPromptFeedback | None = Field(None)
-    usageMetadata: GeminiUsageMetadata | None = Field(None)
-    modelVersion: str | None = Field(None)
+    contents: list[GeminiContent]
+    generationConfig: Optional[GeminiImageGenerationConfig] = None
+    safetySettings: Optional[list[GeminiSafetySetting]] = None
+    systemInstruction: Optional[GeminiSystemInstructionContent] = None
+    tools: Optional[list[GeminiTool]] = None
+    videoMetadata: Optional[GeminiVideoMetadata] = None

comfy_api_nodes/apis/kling_api.py

@@ -1,104 +0,0 @@
-from pydantic import BaseModel, Field
-
-
-class OmniProText2VideoRequest(BaseModel):
-    model_name: str = Field(..., description="kling-video-o1")
-    aspect_ratio: str = Field(..., description="'16:9', '9:16' or '1:1'")
-    duration: str = Field(..., description="'5' or '10'")
-    prompt: str = Field(...)
-    mode: str = Field("pro")
-
-
-class OmniParamImage(BaseModel):
-    image_url: str = Field(...)
-    type: str | None = Field(None, description="Can be 'first_frame' or 'end_frame'")
-
-
-class OmniParamVideo(BaseModel):
-    video_url: str = Field(...)
-    refer_type: str | None = Field(..., description="Can be 'base' or 'feature'")
-    keep_original_sound: str = Field(..., description="'yes' or 'no'")
-
-
-class OmniProFirstLastFrameRequest(BaseModel):
-    model_name: str = Field(..., description="kling-video-o1")
-    image_list: list[OmniParamImage] = Field(..., min_length=1, max_length=7)
-    duration: str = Field(..., description="'5' or '10'")
-    prompt: str = Field(...)
-    mode: str = Field("pro")
-
-
-class OmniProReferences2VideoRequest(BaseModel):
-    model_name: str = Field(..., description="kling-video-o1")
-    aspect_ratio: str | None = Field(..., description="'16:9', '9:16' or '1:1'")
-    image_list: list[OmniParamImage] | None = Field(
-        None, max_length=7, description="Max length 4 when video is present."
-    )
-    video_list: list[OmniParamVideo] | None = Field(None, max_length=1)
-    duration: str | None = Field(..., description="From 3 to 10.")
-    prompt: str = Field(...)
-    mode: str = Field("pro")
-
-
-class TaskStatusVideoResult(BaseModel):
-    duration: str | None = Field(None, description="Total video duration")
-    id: str | None = Field(None, description="Generated video ID")
-    url: str | None = Field(None, description="URL for generated video")
-
-
-class TaskStatusImageResult(BaseModel):
-    index: int = Field(..., description="Image Number，0-9")
-    url: str = Field(..., description="URL for generated image")
-
-
-class TaskStatusResults(BaseModel):
-    videos: list[TaskStatusVideoResult] | None = Field(None)
-    images: list[TaskStatusImageResult] | None = Field(None)
-
-
-class TaskStatusResponseData(BaseModel):
-    created_at: int | None = Field(None, description="Task creation time")
-    updated_at: int | None = Field(None, description="Task update time")
-    task_status: str | None = None
-    task_status_msg: str | None = Field(None, description="Additional failure reason. Only for polling endpoint.")
-    task_id: str | None = Field(None, description="Task ID")
-    task_result: TaskStatusResults | None = Field(None)
-
-
-class TaskStatusResponse(BaseModel):
-    code: int | None = Field(None, description="Error code")
-    message: str | None = Field(None, description="Error message")
-    request_id: str | None = Field(None, description="Request ID")
-    data: TaskStatusResponseData | None = Field(None)
-
-
-class OmniImageParamImage(BaseModel):
-    image: str = Field(...)
-
-
-class OmniProImageRequest(BaseModel):
-    model_name: str = Field(..., description="kling-image-o1")
-    resolution: str = Field(..., description="'1k' or '2k'")
-    aspect_ratio: str | None = Field(...)
-    prompt: str = Field(...)
-    mode: str = Field("pro")
-    n: int | None = Field(1, le=9)
-    image_list: list[OmniImageParamImage] | None = Field(..., max_length=10)
-
-
-class TextToVideoWithAudioRequest(BaseModel):
-    model_name: str = Field(..., description="kling-v2-6")
-    aspect_ratio: str = Field(..., description="'16:9', '9:16' or '1:1'")
-    duration: str = Field(..., description="'5' or '10'")
-    prompt: str = Field(...)
-    mode: str = Field("pro")
-    sound: str = Field(..., description="'on' or 'off'")
-
-
-class ImageToVideoWithAudioRequest(BaseModel):
-    model_name: str = Field(..., description="kling-v2-6")
-    image: str = Field(...)
-    duration: str = Field(..., description="'5' or '10'")
-    prompt: str = Field(...)
-    mode: str = Field("pro")
-    sound: str = Field(..., description="'on' or 'off'")

comfy_api_nodes/apis/openai_api.py

@@ -1,52 +0,0 @@
-from pydantic import BaseModel, Field
-
-
-class Datum2(BaseModel):
-    b64_json: str | None = Field(None, description="Base64 encoded image data")
-    revised_prompt: str | None = Field(None, description="Revised prompt")
-    url: str | None = Field(None, description="URL of the image")
-
-
-class InputTokensDetails(BaseModel):
-    image_tokens: int | None = None
-    text_tokens: int | None = None
-
-
-class Usage(BaseModel):
-    input_tokens: int | None = None
-    input_tokens_details: InputTokensDetails | None = None
-    output_tokens: int | None = None
-    total_tokens: int | None = None
-
-
-class OpenAIImageGenerationResponse(BaseModel):
-    data: list[Datum2] | None = None
-    usage: Usage | None = None
-
-
-class OpenAIImageEditRequest(BaseModel):
-    background: str | None = Field(None, description="Background transparency")
-    model: str = Field(...)
-    moderation: str | None = Field(None)
-    n: int | None = Field(None, description="The number of images to generate")
-    output_compression: int | None = Field(None, description="Compression level for JPEG or WebP (0-100)")
-    output_format: str | None = Field(None)
-    prompt: str = Field(...)
-    quality: str | None = Field(None, description="Size of the image (e.g., 1024x1024, 1536x1024, auto)")
-    size: str | None = Field(None, description="Size of the output image")
-
-
-class OpenAIImageGenerationRequest(BaseModel):
-    background: str | None = Field(None, description="Background transparency")
-    model: str | None = Field(None)
-    moderation: str | None = Field(None)
-    n: int | None = Field(
-        None,
-        description="The number of images to generate.",
-    )
-    output_compression: int | None = Field(None, description="Compression level for JPEG or WebP (0-100)")
-    output_format: str | None = Field(None)
-    prompt: str = Field(...)
-    quality: str | None = Field(None, description="The quality of the generated image")
-    size: str | None = Field(None, description="Size of the image (e.g., 1024x1024, 1536x1024, auto)")
-    style: str | None = Field(None, description="Style of the image (only for dall-e-3)")

comfy_api_nodes/apis/pika_api.py

@@ -0,0 +1,100 @@
+from typing import Optional
+from enum import Enum
+from pydantic import BaseModel, Field
+
+
+class Pikaffect(str, Enum):
+    Cake_ify = "Cake-ify"
+    Crumble = "Crumble"
+    Crush = "Crush"
+    Decapitate = "Decapitate"
+    Deflate = "Deflate"
+    Dissolve = "Dissolve"
+    Explode = "Explode"
+    Eye_pop = "Eye-pop"
+    Inflate = "Inflate"
+    Levitate = "Levitate"
+    Melt = "Melt"
+    Peel = "Peel"
+    Poke = "Poke"
+    Squish = "Squish"
+    Ta_da = "Ta-da"
+    Tear = "Tear"
+
+
+class PikaBodyGenerate22C2vGenerate22PikascenesPost(BaseModel):
+    aspectRatio: Optional[float] = Field(None, description='Aspect ratio (width / height)')
+    duration: Optional[int] = Field(5)
+    ingredientsMode: str = Field(...)
+    negativePrompt: Optional[str] = Field(None)
+    promptText: Optional[str] = Field(None)
+    resolution: Optional[str] = Field('1080p')
+    seed: Optional[int] = Field(None)
+
+
+class PikaGenerateResponse(BaseModel):
+    video_id: str = Field(...)
+
+
+class PikaBodyGenerate22I2vGenerate22I2vPost(BaseModel):
+    duration: Optional[int] = 5
+    negativePrompt: Optional[str] = Field(None)
+    promptText: Optional[str] = Field(None)
+    resolution: Optional[str] = '1080p'
+    seed: Optional[int] = Field(None)
+
+
+class PikaBodyGenerate22KeyframeGenerate22PikaframesPost(BaseModel):
+    duration: Optional[int] = Field(None, ge=5, le=10)
+    negativePrompt: Optional[str] = Field(None)
+    promptText: str = Field(...)
+    resolution: Optional[str] = '1080p'
+    seed: Optional[int] = Field(None)
+
+
+class PikaBodyGenerate22T2vGenerate22T2vPost(BaseModel):
+    aspectRatio: Optional[float] = Field(
+        1.7777777777777777,
+        description='Aspect ratio (width / height)',
+        ge=0.4,
+        le=2.5,
+    )
+    duration: Optional[int] = 5
+    negativePrompt: Optional[str] = Field(None)
+    promptText: str = Field(...)
+    resolution: Optional[str] = '1080p'
+    seed: Optional[int] = Field(None)
+
+
+class PikaBodyGeneratePikadditionsGeneratePikadditionsPost(BaseModel):
+    negativePrompt: Optional[str] = Field(None)
+    promptText: Optional[str] = Field(None)
+    seed: Optional[int] = Field(None)
+
+
+class PikaBodyGeneratePikaffectsGeneratePikaffectsPost(BaseModel):
+    negativePrompt: Optional[str] = Field(None)
+    pikaffect: Optional[str] = None
+    promptText: Optional[str] = Field(None)
+    seed: Optional[int] = Field(None)
+
+
+class PikaBodyGeneratePikaswapsGeneratePikaswapsPost(BaseModel):
+    negativePrompt: Optional[str] = Field(None)
+    promptText: Optional[str] = Field(None)
+    seed: Optional[int] = Field(None)
+    modifyRegionRoi: Optional[str] = Field(None)
+
+
+class PikaStatusEnum(str, Enum):
+    queued = "queued"
+    started = "started"
+    finished = "finished"
+    failed = "failed"
+
+
+class PikaVideoResponse(BaseModel):
+    id: str = Field(...)
+    progress: Optional[int] = Field(None)
+    status: PikaStatusEnum
+    url: Optional[str] = Field(None)

comfy_api_nodes/apis/topaz_api.py

@@ -1,133 +0,0 @@
-from typing import Optional, Union
-
-from pydantic import BaseModel, Field
-
-
-class ImageEnhanceRequest(BaseModel):
-    model: str = Field("Reimagine")
-    output_format: str = Field("jpeg")
-    subject_detection: str = Field("All")
-    face_enhancement: bool = Field(True)
-    face_enhancement_creativity: float = Field(0, description="Is ignored if face_enhancement is false")
-    face_enhancement_strength: float = Field(0.8, description="Is ignored if face_enhancement is false")
-    source_url: str = Field(...)
-    output_width: Optional[int] = Field(None)
-    output_height: Optional[int] = Field(None)
-    crop_to_fill: bool = Field(False)
-    prompt: Optional[str] = Field(None, description="Text prompt for creative upscaling guidance")
-    creativity: int = Field(3, description="Creativity settings range from 1 to 9")
-    face_preservation: str = Field("true", description="To preserve the identity of characters")
-    color_preservation: str = Field("true", description="To preserve the original color")
-
-
-class ImageAsyncTaskResponse(BaseModel):
-    process_id: str = Field(...)
-
-
-class ImageStatusResponse(BaseModel):
-    process_id: str = Field(...)
-    status: str = Field(...)
-    progress: Optional[int] = Field(None)
-    credits: int = Field(...)
-
-
-class ImageDownloadResponse(BaseModel):
-    download_url: str = Field(...)
-    expiry: int = Field(...)
-
-
-class Resolution(BaseModel):
-    width: int = Field(...)
-    height: int = Field(...)
-
-
-class CreateCreateVideoRequestSource(BaseModel):
-    container: str = Field(...)
-    size: int = Field(..., description="Size of the video file in bytes")
-    duration: int = Field(..., description="Duration of the video file in seconds")
-    frameCount: int = Field(..., description="Total number of frames in the video")
-    frameRate: int = Field(...)
-    resolution: Resolution = Field(...)
-
-
-class VideoFrameInterpolationFilter(BaseModel):
-    model: str = Field(...)
-    slowmo: Optional[int] = Field(None)
-    fps: int = Field(...)
-    duplicate: bool = Field(...)
-    duplicate_threshold: float = Field(...)
-
-
-class VideoEnhancementFilter(BaseModel):
-    model: str = Field(...)
-    auto: Optional[str] = Field(None, description="Auto, Manual, Relative")
-    focusFixLevel: Optional[str] = Field(None, description="Downscales video input for correction of blurred subjects")
-    compression: Optional[float] = Field(None, description="Strength of compression recovery")
-    details: Optional[float] = Field(None, description="Amount of detail reconstruction")
-    prenoise: Optional[float] = Field(None, description="Amount of noise to add to input to reduce over-smoothing")
-    noise: Optional[float] = Field(None, description="Amount of noise reduction")
-    halo: Optional[float] = Field(None, description="Amount of halo reduction")
-    preblur: Optional[float] = Field(None, description="Anti-aliasing and deblurring strength")
-    blur: Optional[float] = Field(None, description="Amount of sharpness applied")
-    grain: Optional[float] = Field(None, description="Grain after AI model processing")
-    grainSize: Optional[float] = Field(None, description="Size of generated grain")
-    recoverOriginalDetailValue: Optional[float] = Field(None, description="Source details into the output video")
-    creativity: Optional[str] = Field(None, description="Creativity level(high, low) for slc-1 only")
-    isOptimizedMode: Optional[bool] = Field(None, description="Set to true for Starlight Creative (slc-1) only")
-
-
-class OutputInformationVideo(BaseModel):
-    resolution: Resolution = Field(...)
-    frameRate: int = Field(...)
-    audioCodec: Optional[str] = Field(..., description="Required if audioTransfer is Copy or Convert")
-    audioTransfer: str = Field(..., description="Copy, Convert, None")
-    dynamicCompressionLevel: str = Field(..., description="Low, Mid, High")
-
-
-class Overrides(BaseModel):
-    isPaidDiffusion: bool = Field(True)
-
-
-class CreateVideoRequest(BaseModel):
-    source: CreateCreateVideoRequestSource = Field(...)
-    filters: list[Union[VideoFrameInterpolationFilter, VideoEnhancementFilter]] = Field(...)
-    output: OutputInformationVideo = Field(...)
-    overrides: Overrides = Field(Overrides(isPaidDiffusion=True))
-
-
-class CreateVideoResponse(BaseModel):
-    requestId: str = Field(...)
-
-
-class VideoAcceptResponse(BaseModel):
-    uploadId: str = Field(...)
-    urls: list[str] = Field(...)
-
-
-class VideoCompleteUploadRequestPart(BaseModel):
-    partNum: int = Field(...)
-    eTag: str = Field(...)
-
-
-class VideoCompleteUploadRequest(BaseModel):
-    uploadResults: list[VideoCompleteUploadRequestPart] = Field(...)
-
-
-class VideoCompleteUploadResponse(BaseModel):
-    message: str = Field(..., description="Confirmation message")
-
-
-class VideoStatusResponseEstimates(BaseModel):
-    cost: list[int] = Field(...)
-
-
-class VideoStatusResponseDownloadUrl(BaseModel):
-    url: str = Field(...)
-
-
-class VideoStatusResponse(BaseModel):
-    status: str = Field(...)
-    estimates: Optional[VideoStatusResponseEstimates] = Field(None)
-    progress: Optional[float] = Field(None)
-    message: Optional[str] = Field("")
-    download: Optional[VideoStatusResponseDownloadUrl] = Field(None)

comfy_api_nodes/apis/tripo_api.py

@@ -5,17 +5,11 @@ from typing import Optional, List, Dict, Any, Union
 from pydantic import BaseModel, Field, RootModel
 
 class TripoModelVersion(str, Enum):
-    v3_0_20250812 = 'v3.0-20250812'
     v2_5_20250123 = 'v2.5-20250123'
     v2_0_20240919 = 'v2.0-20240919'
     v1_4_20240625 = 'v1.4-20240625'
 
 
-class TripoGeometryQuality(str, Enum):
-    standard = 'standard'
-    detailed = 'detailed'
-
-
 class TripoTextureQuality(str, Enum):
     standard = 'standard'
     detailed = 'detailed'
@@ -67,20 +61,14 @@ class TripoSpec(str, Enum):
 class TripoAnimation(str, Enum):
     IDLE = "preset:idle"
     WALK = "preset:walk"
-    RUN = "preset:run"
-    DIVE = "preset:dive"
     CLIMB = "preset:climb"
     JUMP = "preset:jump"
+    RUN = "preset:run"
     SLASH = "preset:slash"
     SHOOT = "preset:shoot"
     HURT = "preset:hurt"
     FALL = "preset:fall"
     TURN = "preset:turn"
-    QUADRUPED_WALK = "preset:quadruped:walk"
-    HEXAPOD_WALK = "preset:hexapod:walk"
-    OCTOPOD_WALK = "preset:octopod:walk"
-    SERPENTINE_MARCH = "preset:serpentine:march"
-    AQUATIC_MARCH = "preset:aquatic:march"
 
 class TripoStylizeStyle(str, Enum):
     LEGO = "lego"
@@ -117,11 +105,6 @@ class TripoTaskStatus(str, Enum):
     BANNED = "banned"
     EXPIRED = "expired"
 
-class TripoFbxPreset(str, Enum):
-    BLENDER = "blender"
-    MIXAMO = "mixamo"
-    _3DSMAX = "3dsmax"
-
 class TripoFileTokenReference(BaseModel):
     type: Optional[str] = Field(None, description='The type of the reference')
     file_token: str
@@ -159,7 +142,6 @@ class TripoTextToModelRequest(BaseModel):
     model_seed: Optional[int] = Field(None, description='The seed for the model')
     texture_seed: Optional[int] = Field(None, description='The seed for the texture')
     texture_quality: Optional[TripoTextureQuality] = TripoTextureQuality.standard
-    geometry_quality: Optional[TripoGeometryQuality] = TripoGeometryQuality.standard
     style: Optional[TripoStyle] = None
     auto_size: Optional[bool] = Field(False, description='Whether to auto-size the model')
     quad: Optional[bool] = Field(False, description='Whether to apply quad to the generated model')
@@ -174,7 +156,6 @@ class TripoImageToModelRequest(BaseModel):
     model_seed: Optional[int] = Field(None, description='The seed for the model')
     texture_seed: Optional[int] = Field(None, description='The seed for the texture')
     texture_quality: Optional[TripoTextureQuality] = TripoTextureQuality.standard
-    geometry_quality: Optional[TripoGeometryQuality] = TripoGeometryQuality.standard
     texture_alignment: Optional[TripoTextureAlignment] = Field(TripoTextureAlignment.ORIGINAL_IMAGE, description='The texture alignment method')
     style: Optional[TripoStyle] = Field(None, description='The style to apply to the generated model')
     auto_size: Optional[bool] = Field(False, description='Whether to auto-size the model')
@@ -192,7 +173,6 @@ class TripoMultiviewToModelRequest(BaseModel):
     model_seed: Optional[int] = Field(None, description='The seed for the model')
     texture_seed: Optional[int] = Field(None, description='The seed for the texture')
     texture_quality: Optional[TripoTextureQuality] = TripoTextureQuality.standard
-    geometry_quality: Optional[TripoGeometryQuality] = TripoGeometryQuality.standard
     texture_alignment: Optional[TripoTextureAlignment] = TripoTextureAlignment.ORIGINAL_IMAGE
     auto_size: Optional[bool] = Field(False, description='Whether to auto-size the model')
     orientation: Optional[TripoOrientation] = Field(TripoOrientation.DEFAULT, description='The orientation for the model')
@@ -239,24 +219,14 @@ class TripoConvertModelRequest(BaseModel):
     type: TripoTaskType = Field(TripoTaskType.CONVERT_MODEL, description='Type of task')
     format: TripoConvertFormat = Field(..., description='The format to convert to')
     original_model_task_id: str = Field(..., description='The task ID of the original model')
-    quad: Optional[bool] = Field(None, description='Whether to apply quad to the model')
-    force_symmetry: Optional[bool] = Field(None, description='Whether to force symmetry')
-    face_limit: Optional[int] = Field(None, description='The number of faces to limit the conversion to')
-    flatten_bottom: Optional[bool] = Field(None, description='Whether to flatten the bottom of the model')
-    flatten_bottom_threshold: Optional[float] = Field(None, description='The threshold for flattening the bottom')
-    texture_size: Optional[int] = Field(None, description='The size of the texture')
+    quad: Optional[bool] = Field(False, description='Whether to apply quad to the model')
+    force_symmetry: Optional[bool] = Field(False, description='Whether to force symmetry')
+    face_limit: Optional[int] = Field(10000, description='The number of faces to limit the conversion to')
+    flatten_bottom: Optional[bool] = Field(False, description='Whether to flatten the bottom of the model')
+    flatten_bottom_threshold: Optional[float] = Field(0.01, description='The threshold for flattening the bottom')
+    texture_size: Optional[int] = Field(4096, description='The size of the texture')
     texture_format: Optional[TripoTextureFormat] = Field(TripoTextureFormat.JPEG, description='The format of the texture')
-    pivot_to_center_bottom: Optional[bool] = Field(None, description='Whether to pivot to the center bottom')
-    scale_factor: Optional[float] = Field(None, description='The scale factor for the model')
-    with_animation: Optional[bool] = Field(None, description='Whether to include animations')
-    pack_uv: Optional[bool] = Field(None, description='Whether to pack the UVs')
-    bake: Optional[bool] = Field(None, description='Whether to bake the model')
-    part_names: Optional[List[str]] = Field(None, description='The names of the parts to include')
-    fbx_preset: Optional[TripoFbxPreset] = Field(None, description='The preset for the FBX export')
-    export_vertex_colors: Optional[bool] = Field(None, description='Whether to export the vertex colors')
-    export_orientation: Optional[TripoOrientation] = Field(None, description='The orientation for the export')
-    animate_in_place: Optional[bool] = Field(None, description='Whether to animate in place')
-
+    pivot_to_center_bottom: Optional[bool] = Field(False, description='Whether to pivot to the center bottom')
 
 class TripoTaskRequest(RootModel):
     root: Union[