Merge remote-tracking branch 'upstream/master'

# Conflicts:
#	.gitignore

.ci/update_windows/update.py

@@ -75,6 +75,25 @@ else:
 print("pulling latest changes")
 pull(repo)
 
+if "--stable" in sys.argv:
+    def latest_tag(repo):
+        versions = []
+        for k in repo.references:
+            try:
+                prefix = "refs/tags/v"
+                if k.startswith(prefix):
+                    version = list(map(int, k[len(prefix):].split(".")))
+                    versions.append((version[0] * 10000000000 + version[1] * 100000 + version[2], k))
+            except:
+                pass
+        versions.sort()
+        if len(versions) > 0:
+            return versions[-1][1]
+        return None
+    latest_tag = latest_tag(repo)
+    if latest_tag is not None:
+        repo.checkout(latest_tag)
+
 print("Done!")
 
 self_update = True
@@ -115,3 +134,13 @@ if not os.path.exists(req_path) or not files_equal(repo_req_path, req_path):
         shutil.copy(repo_req_path, req_path)
     except:
         pass
+
+
+stable_update_script = os.path.join(repo_path, ".ci/update_windows/update_comfyui_stable.bat")
+stable_update_script_to = os.path.join(cur_path, "update_comfyui_stable.bat")
+
+try:
+    if not file_size(stable_update_script_to) > 10:
+        shutil.copy(stable_update_script, stable_update_script_to)
+except:
+    pass

.ci/update_windows/update_comfyui_stable.bat

@@ -0,0 +1,8 @@
+@echo off
+..\python_embeded\python.exe .\update.py ..\ComfyUI\ --stable
+if exist update_new.py (
+  move /y update_new.py update.py
+  echo Running updater again since it got updated.
+  ..\python_embeded\python.exe .\update.py ..\ComfyUI\ --skip_self_update --stable
+)
+if "%~1"=="" pause

.gitattributes

@@ -0,0 +1,2 @@
+/web/assets/** linguist-generated
+/web/** linguist-vendored

.github/workflows/pullrequest-ci-run.yml

@@ -0,0 +1,53 @@
+# This is the GitHub Workflow that drives full-GPU-enabled tests of pull requests to ComfyUI, when the 'Run-CI-Test' label is added
+# Results are reported as checkmarks on the commits, as well as onto https://ci.comfy.org/
+name: Pull Request CI Workflow Runs
+on:
+    pull_request_target:
+        types: [labeled]
+
+jobs:
+  pr-test-stable:
+    if: ${{ github.event.label.name == 'Run-CI-Test' }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [macos, linux, windows]
+        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: linux
+            runner_label: [self-hosted, Linux]
+            flags: ""
+          - os: windows
+            runner_label: [self-hosted, win]
+            flags: ""
+    runs-on: ${{ matrix.runner_label }}
+    steps:
+      - name: Test Workflows
+        uses: comfy-org/comfy-action@main
+        with:
+          os: ${{ matrix.os }}
+          python_version: ${{ matrix.python_version }}
+          torch_version: ${{ matrix.torch_version }}
+          google_credentials: ${{ secrets.GCS_SERVICE_ACCOUNT_JSON }}
+          comfyui_flags: ${{ matrix.flags }}
+          use_prior_commit: 'true'
+  comment:
+    if: ${{ github.event.label.name == 'Run-CI-Test' }}
+    runs-on: ubuntu-latest
+    permissions:
+      pull-requests: write
+    steps:
+      - uses: actions/github-script@v6
+        with:
+          script: |
+            github.rest.issues.createComment({
+              issue_number: context.issue.number,
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              body: '(Automated Bot Message) CI Tests are running, you can view the results at https://ci.comfy.org/?branch=${{ github.event.pull_request.number }}%2Fmerge'
+            })

.github/workflows/test-ci.yml

@@ -0,0 +1,95 @@
+# This is the GitHub Workflow that drives automatic full-GPU-enabled tests of all new commits to the master branch of ComfyUI
+# Results are reported as checkmarks on the commits, as well as onto https://ci.comfy.org/
+name: Full Comfy CI Workflow Runs
+on:
+  push:
+    branches:
+      - master
+    paths-ignore:
+      - 'app/**'
+      - 'input/**'
+      - 'output/**'
+      - 'notebooks/**'
+      - 'script_examples/**'
+      - '.github/**'
+      - 'web/**'
+  workflow_dispatch:
+
+jobs:
+  test-stable:
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [macos, linux, windows]
+        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: linux
+            runner_label: [self-hosted, Linux]
+            flags: ""
+          - os: windows
+            runner_label: [self-hosted, win]
+            flags: ""
+    runs-on: ${{ matrix.runner_label }}
+    steps:
+      - name: Test Workflows
+        uses: comfy-org/comfy-action@main
+        with:
+          os: ${{ matrix.os }}
+          python_version: ${{ matrix.python_version }}
+          torch_version: ${{ matrix.torch_version }}
+          google_credentials: ${{ secrets.GCS_SERVICE_ACCOUNT_JSON }}
+          comfyui_flags: ${{ matrix.flags }}
+
+  test-win-nightly:
+    strategy:
+      fail-fast: true
+      matrix:
+        os: [windows]
+        python_version: ["3.9", "3.10", "3.11", "3.12"]
+        cuda_version: ["12.1"]
+        torch_version: ["nightly"]
+        include:
+          - os: windows
+            runner_label: [self-hosted, win]
+            flags: ""
+    runs-on: ${{ matrix.runner_label }}
+    steps:
+      - name: Test Workflows
+        uses: comfy-org/comfy-action@main
+        with:
+          os: ${{ matrix.os }}
+          python_version: ${{ matrix.python_version }}
+          torch_version: ${{ matrix.torch_version }}
+          google_credentials: ${{ secrets.GCS_SERVICE_ACCOUNT_JSON }}
+          comfyui_flags: ${{ matrix.flags }}
+
+  test-unix-nightly:
+    strategy:
+      fail-fast: false
+      matrix:
+        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: linux
+            runner_label: [self-hosted, Linux]
+            flags: ""
+    runs-on: ${{ matrix.runner_label }}
+    steps:
+      - name: Test Workflows
+        uses: comfy-org/comfy-action@main
+        with:
+          os: ${{ matrix.os }}
+          python_version: ${{ matrix.python_version }}
+          torch_version: ${{ matrix.torch_version }}
+          google_credentials: ${{ secrets.GCS_SERVICE_ACCOUNT_JSON }}
+          comfyui_flags: ${{ matrix.flags }}

.github/workflows/test-launch.yml

@@ -1,10 +1,4 @@
-# This is a temporary action during frontend TS migration.
-# This file should be removed after TS migration is completed.
-# The browser test is here to ensure TS repo is working the same way as the
-# current JS code.
-# If you are adding UI feature, please sync your changes to the TS repo:
-# huchenlei/ComfyUI_frontend and update test expectation files accordingly.
-name: Playwright Browser Tests CI
+name: Test server launches without errors
 
 on:
   push:
@@ -21,15 +15,6 @@ jobs:
       with:
         repository: "comfyanonymous/ComfyUI"
         path: "ComfyUI"
-    - name: Checkout ComfyUI_frontend
-      uses: actions/checkout@v4
-      with:
-        repository: "huchenlei/ComfyUI_frontend"
-        path: "ComfyUI_frontend"
-        ref: "fcc54d803e5b6a9b08a462a1d94899318c96dcbb"
-    - uses: actions/setup-node@v3
-      with:
-        node-version: lts/*
     - uses: actions/setup-python@v4
       with:
         python-version: '3.8'
@@ -45,16 +30,6 @@ jobs:
         python main.py --cpu 2>&1 | tee console_output.log &
         wait-for-it --service 127.0.0.1:8188 -t 600
       working-directory: ComfyUI
-    - name: Install ComfyUI_frontend dependencies
-      run: |
-        npm ci
-      working-directory: ComfyUI_frontend
-    - name: Install Playwright Browsers
-      run: npx playwright install --with-deps
-      working-directory: ComfyUI_frontend
-    - name: Run Playwright tests
-      run: npx playwright test
-      working-directory: ComfyUI_frontend
     - name: Check for unhandled exceptions in server log
       run: |
         if grep -qE "Exception|Error" console_output.log; then
@@ -62,12 +37,6 @@ jobs:
           exit 1
         fi
       working-directory: ComfyUI
-    - uses: actions/upload-artifact@v4
-      if: always()
-      with:
-        name: playwright-report
-        path: ComfyUI_frontend/playwright-report/
-        retention-days: 30
     - uses: actions/upload-artifact@v4
       if: always()
       with:

.github/workflows/test-ui.yaml

@@ -1,30 +0,0 @@
-name: Tests CI
-
-on: [push, pull_request]
-
-jobs:
-  test:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v4
-    - uses: actions/setup-node@v3
-      with:
-        node-version: 18
-    - uses: actions/setup-python@v4
-      with: 
-        python-version: '3.10'
-    - name: Install requirements
-      run: |
-        python -m pip install --upgrade pip
-        pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
-        pip install -r requirements.txt
-    - name: Run Tests
-      run: | 
-        npm ci
-        npm run test:generate
-        npm test -- --verbose
-      working-directory: ./tests-ui
-    - name: Run Unit Tests
-      run: |
-        pip install -r tests-unit/requirements.txt
-        python -m pytest tests-unit

.gitignore

@@ -21,5 +21,5 @@ desktop_shortcut
 comfyui.log
 comfyui.prev.log
 *.log
-
 web_custom_versions/
+.DS_Store

comfy/cli_args.py

@@ -92,6 +92,10 @@ class LatentPreviewMethod(enum.Enum):
 
 parser.add_argument("--preview-method", type=LatentPreviewMethod, default=LatentPreviewMethod.NoPreviews, help="Default preview method for sampler nodes.", action=EnumAction)
 
+cache_group = parser.add_mutually_exclusive_group()
+cache_group.add_argument("--cache-classic", action="store_true", help="Use the old style (aggressive) caching.")
+cache_group.add_argument("--cache-lru", type=int, default=0, help="Use LRU caching with a maximum of N node results cached. May use more RAM/VRAM.")
+
 attn_group = parser.add_mutually_exclusive_group()
 attn_group.add_argument("--use-split-cross-attention", action="store_true", help="Use the split cross attention optimization. Ignored when xformers is used.")
 attn_group.add_argument("--use-quad-cross-attention", action="store_true", help="Use the sub-quadratic cross attention optimization . Ignored when xformers is used.")

comfy/controlnet.py

@@ -1,4 +1,24 @@
+"""
+    This file is part of ComfyUI.
+    Copyright (C) 2024 Comfy
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+
 import torch
+from enum import Enum
 import math
 import os
 import logging
@@ -13,6 +33,8 @@ import comfy.cldm.cldm
 import comfy.t2i_adapter.adapter
 import comfy.ldm.cascade.controlnet
 import comfy.cldm.mmdit
+import comfy.ldm.hydit.controlnet
+import comfy.ldm.flux.controlnet_xlabs
 
 
 def broadcast_image_to(tensor, target_batch_size, batched_number):
@@ -33,6 +55,10 @@ def broadcast_image_to(tensor, target_batch_size, batched_number):
     else:
         return torch.cat([tensor] * batched_number, dim=0)
 
+class StrengthType(Enum):
+    CONSTANT = 1
+    LINEAR_UP = 2
+
 class ControlBase:
     def __init__(self, device=None):
         self.cond_hint_original = None
@@ -51,6 +77,8 @@ class ControlBase:
             device = comfy.model_management.get_torch_device()
         self.device = device
         self.previous_controlnet = None
+        self.extra_conds = []
+        self.strength_type = StrengthType.CONSTANT
 
     def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None):
         self.cond_hint_original = cond_hint
@@ -93,6 +121,8 @@ class ControlBase:
         c.latent_format = self.latent_format
         c.extra_args = self.extra_args.copy()
         c.vae = self.vae
+        c.extra_conds = self.extra_conds.copy()
+        c.strength_type = self.strength_type
 
     def inference_memory_requirements(self, dtype):
         if self.previous_controlnet is not None:
@@ -113,7 +143,10 @@ class ControlBase:
 
                     if x not in applied_to: #memory saving strategy, allow shared tensors and only apply strength to shared tensors once
                         applied_to.add(x)
-                        x *= self.strength
+                        if self.strength_type == StrengthType.CONSTANT:
+                            x *= self.strength
+                        elif self.strength_type == StrengthType.LINEAR_UP:
+                            x *= (self.strength ** float(len(control_output) - i))
 
                     if x.dtype != output_dtype:
                         x = x.to(output_dtype)
@@ -142,7 +175,7 @@ class ControlBase:
 
 
 class ControlNet(ControlBase):
-    def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None):
+    def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT):
         super().__init__(device)
         self.control_model = control_model
         self.load_device = load_device
@@ -154,6 +187,8 @@ class ControlNet(ControlBase):
         self.model_sampling_current = None
         self.manual_cast_dtype = manual_cast_dtype
         self.latent_format = latent_format
+        self.extra_conds += extra_conds
+        self.strength_type = strength_type
 
     def get_control(self, x_noisy, t, cond, batched_number):
         control_prev = None
@@ -191,13 +226,16 @@ class ControlNet(ControlBase):
             self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
 
         context = cond.get('crossattn_controlnet', cond['c_crossattn'])
-        y = cond.get('y', None)
-        if y is not None:
-            y = y.to(dtype)
+        extra = self.extra_args.copy()
+        for c in self.extra_conds:
+            temp = cond.get(c, None)
+            if temp is not None:
+                extra[c] = temp.to(dtype)
+
         timestep = self.model_sampling_current.timestep(t)
         x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
 
-        control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y, **self.extra_args)
+        control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.to(dtype), context=context.to(dtype), **extra)
         return self.control_merge(control, control_prev, output_dtype)
 
     def copy(self):
@@ -286,6 +324,7 @@ class ControlLora(ControlNet):
         ControlBase.__init__(self, device)
         self.control_weights = control_weights
         self.global_average_pooling = global_average_pooling
+        self.extra_conds += ["y"]
 
     def pre_run(self, model, percent_to_timestep_function):
         super().pre_run(model, percent_to_timestep_function)
@@ -338,12 +377,8 @@ class ControlLora(ControlNet):
     def inference_memory_requirements(self, dtype):
         return comfy.utils.calculate_parameters(self.control_weights) * comfy.model_management.dtype_size(dtype) + ControlBase.inference_memory_requirements(self, dtype)
 
-def load_controlnet_mmdit(sd):
-    new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "")
-    model_config = comfy.model_detection.model_config_from_unet(new_sd, "", True)
-    num_blocks = comfy.model_detection.count_blocks(new_sd, 'joint_blocks.{}.')
-    for k in sd:
-        new_sd[k] = sd[k]
+def controlnet_config(sd):
+    model_config = comfy.model_detection.model_config_from_unet(sd, "", True)
 
     supported_inference_dtypes = model_config.supported_inference_dtypes
 
@@ -356,14 +391,27 @@ def load_controlnet_mmdit(sd):
     else:
         operations = comfy.ops.disable_weight_init
 
-    control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, operations=operations, device=load_device, dtype=unet_dtype, **controlnet_config)
-    missing, unexpected = control_model.load_state_dict(new_sd, strict=False)
+    return model_config, operations, load_device, unet_dtype, manual_cast_dtype
+
+def controlnet_load_state_dict(control_model, sd):
+    missing, unexpected = control_model.load_state_dict(sd, strict=False)
 
     if len(missing) > 0:
         logging.warning("missing controlnet keys: {}".format(missing))
 
     if len(unexpected) > 0:
         logging.debug("unexpected controlnet keys: {}".format(unexpected))
+    return control_model
+
+def load_controlnet_mmdit(sd):
+    new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "")
+    model_config, operations, load_device, unet_dtype, manual_cast_dtype = controlnet_config(new_sd)
+    num_blocks = comfy.model_detection.count_blocks(new_sd, 'joint_blocks.{}.')
+    for k in sd:
+        new_sd[k] = sd[k]
+
+    control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, operations=operations, device=load_device, dtype=unet_dtype, **model_config.unet_config)
+    control_model = controlnet_load_state_dict(control_model, new_sd)
 
     latent_format = comfy.latent_formats.SD3()
     latent_format.shift_factor = 0 #SD3 controlnet weirdness
@@ -371,8 +419,31 @@ def load_controlnet_mmdit(sd):
     return control
 
 
+def load_controlnet_hunyuandit(controlnet_data):
+    model_config, operations, load_device, unet_dtype, manual_cast_dtype = controlnet_config(controlnet_data)
+
+    control_model = comfy.ldm.hydit.controlnet.HunYuanControlNet(operations=operations, device=load_device, dtype=unet_dtype)
+    control_model = controlnet_load_state_dict(control_model, controlnet_data)
+
+    latent_format = comfy.latent_formats.SDXL()
+    extra_conds = ['text_embedding_mask', 'encoder_hidden_states_t5', 'text_embedding_mask_t5', 'image_meta_size', 'style', 'cos_cis_img', 'sin_cis_img']
+    control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds, strength_type=StrengthType.CONSTANT)
+    return control
+
+def load_controlnet_flux_xlabs(sd):
+    model_config, operations, load_device, unet_dtype, manual_cast_dtype = controlnet_config(sd)
+    control_model = comfy.ldm.flux.controlnet_xlabs.ControlNetFlux(operations=operations, device=load_device, dtype=unet_dtype, **model_config.unet_config)
+    control_model = controlnet_load_state_dict(control_model, sd)
+    extra_conds = ['y', 'guidance']
+    control = ControlNet(control_model, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
+    return control
+
+
 def load_controlnet(ckpt_path, model=None):
     controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
+    if 'after_proj_list.18.bias' in controlnet_data.keys(): #Hunyuan DiT
+        return load_controlnet_hunyuandit(controlnet_data)
+
     if "lora_controlnet" in controlnet_data:
         return ControlLora(controlnet_data)
 
@@ -430,7 +501,10 @@ def load_controlnet(ckpt_path, model=None):
             logging.warning("leftover keys: {}".format(leftover_keys))
         controlnet_data = new_sd
     elif "controlnet_blocks.0.weight" in controlnet_data: #SD3 diffusers format
-        return load_controlnet_mmdit(controlnet_data)
+        if "double_blocks.0.img_attn.norm.key_norm.scale" in controlnet_data:
+            return load_controlnet_flux_xlabs(controlnet_data)
+        else:
+            return load_controlnet_mmdit(controlnet_data)
 
     pth_key = 'control_model.zero_convs.0.0.weight'
     pth = False

comfy/diffusers_load.py

@@ -22,7 +22,7 @@ def load_diffusers(model_path, output_vae=True, output_clip=True, embedding_dire
     if text_encoder2_path is not None:
         text_encoder_paths.append(text_encoder2_path)
 
-    unet = comfy.sd.load_unet(unet_path)
+    unet = comfy.sd.load_diffusion_model(unet_path)
 
     clip = None
     if output_clip:

comfy/ldm/flux/controlnet_xlabs.py

@@ -0,0 +1,104 @@
+#Original code can be found on: https://github.com/XLabs-AI/x-flux/blob/main/src/flux/controlnet.py
+
+import torch
+from torch import Tensor, nn
+from einops import rearrange, repeat
+
+from .layers import (DoubleStreamBlock, EmbedND, LastLayer,
+                                 MLPEmbedder, SingleStreamBlock,
+                                 timestep_embedding)
+
+from .model import Flux
+import comfy.ldm.common_dit
+
+
+class ControlNetFlux(Flux):
+    def __init__(self, image_model=None, dtype=None, device=None, operations=None, **kwargs):
+        super().__init__(final_layer=False, dtype=dtype, device=device, operations=operations, **kwargs)
+
+        # add ControlNet blocks
+        self.controlnet_blocks = nn.ModuleList([])
+        for _ in range(self.params.depth):
+            controlnet_block = operations.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device)
+            # controlnet_block = zero_module(controlnet_block)
+            self.controlnet_blocks.append(controlnet_block)
+        self.pos_embed_input = operations.Linear(self.in_channels, self.hidden_size, bias=True, dtype=dtype, device=device)
+        self.gradient_checkpointing = False
+        self.input_hint_block = nn.Sequential(
+            operations.Conv2d(3, 16, 3, padding=1, dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.Conv2d(16, 16, 3, padding=1, stride=2, dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.Conv2d(16, 16, 3, padding=1, dtype=dtype, device=device)
+        )
+
+    def forward_orig(
+        self,
+        img: Tensor,
+        img_ids: Tensor,
+        controlnet_cond: Tensor,
+        txt: Tensor,
+        txt_ids: Tensor,
+        timesteps: Tensor,
+        y: Tensor,
+        guidance: Tensor = None,
+    ) -> Tensor:
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+        # running on sequences img
+        img = self.img_in(img)
+        controlnet_cond = self.input_hint_block(controlnet_cond)
+        controlnet_cond = rearrange(controlnet_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
+        controlnet_cond = self.pos_embed_input(controlnet_cond)
+        img = img + controlnet_cond
+        vec = self.time_in(timestep_embedding(timesteps, 256))
+        if self.params.guidance_embed:
+            vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+        vec = vec + self.vector_in(y)
+        txt = self.txt_in(txt)
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        pe = self.pe_embedder(ids)
+
+        block_res_samples = ()
+
+        for block in self.double_blocks:
+            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+            block_res_samples = block_res_samples + (img,)
+
+        controlnet_block_res_samples = ()
+        for block_res_sample, controlnet_block in zip(block_res_samples, self.controlnet_blocks):
+            block_res_sample = controlnet_block(block_res_sample)
+            controlnet_block_res_samples = controlnet_block_res_samples + (block_res_sample,)
+
+        return {"output": (controlnet_block_res_samples * 10)[:19]}
+
+    def forward(self, x, timesteps, context, y, guidance=None, hint=None, **kwargs):
+        hint = hint * 2.0 - 1.0
+
+        bs, c, h, w = x.shape
+        patch_size = 2
+        x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
+
+        img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
+
+        h_len = ((h + (patch_size // 2)) // patch_size)
+        w_len = ((w + (patch_size // 2)) // patch_size)
+        img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
+        img_ids[..., 1] = img_ids[..., 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype)[:, None]
+        img_ids[..., 2] = img_ids[..., 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype)[None, :]
+        img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+        txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
+        return self.forward_orig(img, img_ids, hint, context, txt_ids, timesteps, y, guidance)

comfy/ldm/flux/layers.py

@@ -2,12 +2,12 @@ import math
 from dataclasses import dataclass
 
 import torch
-from einops import rearrange
 from torch import Tensor, nn
 
 from .math import attention, rope
 import comfy.ops
 
+
 class EmbedND(nn.Module):
     def __init__(self, dim: int, theta: int, axes_dim: list):
         super().__init__()
@@ -36,9 +36,7 @@ def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 10
     """
     t = time_factor * t
     half = dim // 2
-    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
-        t.device
-    )
+    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half)
 
     args = t[:, None].float() * freqs[None]
     embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
@@ -48,7 +46,6 @@ def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 10
         embedding = embedding.to(t)
     return embedding
 
-
 class MLPEmbedder(nn.Module):
     def __init__(self, in_dim: int, hidden_dim: int, dtype=None, device=None, operations=None):
         super().__init__()
@@ -94,14 +91,6 @@ class SelfAttention(nn.Module):
         self.norm = QKNorm(head_dim, dtype=dtype, device=device, operations=operations)
         self.proj = operations.Linear(dim, dim, dtype=dtype, device=device)
 
-    def forward(self, x: Tensor, pe: Tensor) -> Tensor:
-        qkv = self.qkv(x)
-        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-        q, k = self.norm(q, k, v)
-        x = attention(q, k, v, pe=pe)
-        x = self.proj(x)
-        return x
-
 
 @dataclass
 class ModulationOut:
@@ -163,22 +152,21 @@ class DoubleStreamBlock(nn.Module):
         img_modulated = self.img_norm1(img)
         img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
         img_qkv = self.img_attn.qkv(img_modulated)
-        img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        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 = self.txt_norm1(txt)
         txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
         txt_qkv = self.txt_attn.qkv(txt_modulated)
-        txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        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
-        q = torch.cat((txt_q, img_q), dim=2)
-        k = torch.cat((txt_k, img_k), dim=2)
-        v = torch.cat((txt_v, img_v), dim=2)
+        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)
 
-        attn = attention(q, k, v, pe=pe)
         txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
 
         # calculate the img bloks
@@ -186,8 +174,12 @@ class DoubleStreamBlock(nn.Module):
         img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
 
         # calculate the txt bloks
-        txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
-        txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+        txt += txt_mod1.gate * self.txt_attn.proj(txt_attn)
+        txt += txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+
+        if txt.dtype == torch.float16:
+            txt = txt.clip(-65504, 65504)
+
         return img, txt
 
 
@@ -232,14 +224,17 @@ class SingleStreamBlock(nn.Module):
         x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
         qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
 
-        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        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)
         # compute activation in mlp stream, cat again and run second linear layer
         output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
-        return x + mod.gate * output
+        x += mod.gate * output
+        if x.dtype == torch.float16:
+            x = x.clip(-65504, 65504)
+        return x
 
 
 class LastLayer(nn.Module):

comfy/ldm/flux/model.py

@@ -38,7 +38,7 @@ class Flux(nn.Module):
     Transformer model for flow matching on sequences.
     """
 
-    def __init__(self, image_model=None, dtype=None, device=None, operations=None, **kwargs):
+    def __init__(self, image_model=None, final_layer=True, dtype=None, device=None, operations=None, **kwargs):
         super().__init__()
         self.dtype = dtype
         params = FluxParams(**kwargs)
@@ -83,7 +83,8 @@ class Flux(nn.Module):
             ]
         )
 
-        self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations)
+        if final_layer:
+            self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations)
 
     def forward_orig(
         self,
@@ -94,6 +95,7 @@ class Flux(nn.Module):
         timesteps: Tensor,
         y: Tensor,
         guidance: Tensor = None,
+        control=None,
     ) -> Tensor:
         if img.ndim != 3 or txt.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
@@ -112,8 +114,15 @@ class Flux(nn.Module):
         ids = torch.cat((txt_ids, img_ids), dim=1)
         pe = self.pe_embedder(ids)
 
-        for block in self.double_blocks:
-            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+        for i in range(len(self.double_blocks)):
+            img, txt = self.double_blocks[i](img=img, txt=txt, vec=vec, pe=pe)
+
+            if control is not None: #Controlnet
+                control_o = control.get("output")
+                if i < len(control_o):
+                    add = control_o[i]
+                    if add is not None:
+                        img += add
 
         img = torch.cat((txt, img), 1)
         for block in self.single_blocks:
@@ -123,7 +132,7 @@ class Flux(nn.Module):
         img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
         return img
 
-    def forward(self, x, timestep, context, y, guidance, **kwargs):
+    def forward(self, x, timestep, context, y, guidance, control=None, **kwargs):
         bs, c, h, w = x.shape
         patch_size = 2
         x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
@@ -138,5 +147,5 @@ class Flux(nn.Module):
         img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
 
         txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
-        out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance)
+        out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control)
         return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w]

comfy/ldm/hydit/attn_layers.py

@@ -47,7 +47,7 @@ def reshape_for_broadcast(freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], x
 
 
 def rotate_half(x):
-    x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+    x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
     return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
 
 
@@ -78,10 +78,9 @@ def apply_rotary_emb(
     xk_out = None
     if isinstance(freqs_cis, tuple):
         cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first)    # [S, D]
-        cos, sin = cos.to(xq.device), sin.to(xq.device)
-        xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq)
+        xq_out = (xq * cos + rotate_half(xq) * sin)
         if xk is not None:
-            xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk)
+            xk_out = (xk * cos + rotate_half(xk) * sin)
     else:
         xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))  # [B, S, H, D//2]
         freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device)   # [S, D//2] --> [1, S, 1, D//2]

comfy/ldm/hydit/controlnet.py

@@ -0,0 +1,321 @@
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch.utils import checkpoint
+
+from comfy.ldm.modules.diffusionmodules.mmdit import (
+    Mlp,
+    TimestepEmbedder,
+    PatchEmbed,
+    RMSNorm,
+)
+from comfy.ldm.modules.diffusionmodules.util import timestep_embedding
+from .poolers import AttentionPool
+
+import comfy.latent_formats
+from .models import HunYuanDiTBlock, calc_rope
+
+from .posemb_layers import get_2d_rotary_pos_embed, get_fill_resize_and_crop
+
+
+class HunYuanControlNet(nn.Module):
+    """
+    HunYuanDiT: Diffusion model with a Transformer backbone.
+
+    Inherit ModelMixin and ConfigMixin to be compatible with the sampler StableDiffusionPipeline of diffusers.
+
+    Inherit PeftAdapterMixin to be compatible with the PEFT training pipeline.
+
+    Parameters
+    ----------
+    args: argparse.Namespace
+        The arguments parsed by argparse.
+    input_size: tuple
+        The size of the input image.
+    patch_size: int
+        The size of the patch.
+    in_channels: int
+        The number of input channels.
+    hidden_size: int
+        The hidden size of the transformer backbone.
+    depth: int
+        The number of transformer blocks.
+    num_heads: int
+        The number of attention heads.
+    mlp_ratio: float
+        The ratio of the hidden size of the MLP in the transformer block.
+    log_fn: callable
+        The logging function.
+    """
+
+    def __init__(
+        self,
+        input_size: tuple = 128,
+        patch_size: int = 2,
+        in_channels: int = 4,
+        hidden_size: int = 1408,
+        depth: int = 40,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.3637,
+        text_states_dim=1024,
+        text_states_dim_t5=2048,
+        text_len=77,
+        text_len_t5=256,
+        qk_norm=True,  # See http://arxiv.org/abs/2302.05442 for details.
+        size_cond=False,
+        use_style_cond=False,
+        learn_sigma=True,
+        norm="layer",
+        log_fn: callable = print,
+        attn_precision=None,
+        dtype=None,
+        device=None,
+        operations=None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.log_fn = log_fn
+        self.depth = depth
+        self.learn_sigma = learn_sigma
+        self.in_channels = in_channels
+        self.out_channels = in_channels * 2 if learn_sigma else in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.text_states_dim = text_states_dim
+        self.text_states_dim_t5 = text_states_dim_t5
+        self.text_len = text_len
+        self.text_len_t5 = text_len_t5
+        self.size_cond = size_cond
+        self.use_style_cond = use_style_cond
+        self.norm = norm
+        self.dtype = dtype
+        self.latent_format = comfy.latent_formats.SDXL
+
+        self.mlp_t5 = nn.Sequential(
+            nn.Linear(
+                self.text_states_dim_t5,
+                self.text_states_dim_t5 * 4,
+                bias=True,
+                dtype=dtype,
+                device=device,
+            ),
+            nn.SiLU(),
+            nn.Linear(
+                self.text_states_dim_t5 * 4,
+                self.text_states_dim,
+                bias=True,
+                dtype=dtype,
+                device=device,
+            ),
+        )
+        # learnable replace
+        self.text_embedding_padding = nn.Parameter(
+            torch.randn(
+                self.text_len + self.text_len_t5,
+                self.text_states_dim,
+                dtype=dtype,
+                device=device,
+            )
+        )
+
+        # Attention pooling
+        pooler_out_dim = 1024
+        self.pooler = AttentionPool(
+            self.text_len_t5,
+            self.text_states_dim_t5,
+            num_heads=8,
+            output_dim=pooler_out_dim,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )
+
+        # Dimension of the extra input vectors
+        self.extra_in_dim = pooler_out_dim
+
+        if self.size_cond:
+            # Image size and crop size conditions
+            self.extra_in_dim += 6 * 256
+
+        if self.use_style_cond:
+            # Here we use a default learned embedder layer for future extension.
+            self.style_embedder = nn.Embedding(
+                1, hidden_size, dtype=dtype, device=device
+            )
+            self.extra_in_dim += hidden_size
+
+        # Text embedding for `add`
+        self.x_embedder = PatchEmbed(
+            input_size,
+            patch_size,
+            in_channels,
+            hidden_size,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )
+        self.t_embedder = TimestepEmbedder(
+            hidden_size, dtype=dtype, device=device, operations=operations
+        )
+        self.extra_embedder = nn.Sequential(
+            operations.Linear(
+                self.extra_in_dim, hidden_size * 4, dtype=dtype, device=device
+            ),
+            nn.SiLU(),
+            operations.Linear(
+                hidden_size * 4, hidden_size, bias=True, dtype=dtype, device=device
+            ),
+        )
+
+        # Image embedding
+        num_patches = self.x_embedder.num_patches
+
+        # HUnYuanDiT Blocks
+        self.blocks = nn.ModuleList(
+            [
+                HunYuanDiTBlock(
+                    hidden_size=hidden_size,
+                    c_emb_size=hidden_size,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    text_states_dim=self.text_states_dim,
+                    qk_norm=qk_norm,
+                    norm_type=self.norm,
+                    skip=False,
+                    attn_precision=attn_precision,
+                    dtype=dtype,
+                    device=device,
+                    operations=operations,
+                )
+                for _ in range(19)
+            ]
+        )
+
+        # Input zero linear for the first block
+        self.before_proj = operations.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device)
+
+
+        # Output zero linear for the every block
+        self.after_proj_list = nn.ModuleList(
+            [
+
+                    operations.Linear(
+                        self.hidden_size, self.hidden_size, dtype=dtype, device=device
+                    )
+                for _ in range(len(self.blocks))
+            ]
+        )
+
+    def forward(
+        self,
+        x,
+        hint,
+        timesteps,
+        context,#encoder_hidden_states=None,
+        text_embedding_mask=None,
+        encoder_hidden_states_t5=None,
+        text_embedding_mask_t5=None,
+        image_meta_size=None,
+        style=None,
+        return_dict=False,
+        **kwarg,
+    ):
+        """
+        Forward pass of the encoder.
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            (B, D, H, W)
+        t: torch.Tensor
+            (B)
+        encoder_hidden_states: torch.Tensor
+            CLIP text embedding, (B, L_clip, D)
+        text_embedding_mask: torch.Tensor
+            CLIP text embedding mask, (B, L_clip)
+        encoder_hidden_states_t5: torch.Tensor
+            T5 text embedding, (B, L_t5, D)
+        text_embedding_mask_t5: torch.Tensor
+            T5 text embedding mask, (B, L_t5)
+        image_meta_size: torch.Tensor
+            (B, 6)
+        style: torch.Tensor
+            (B)
+        cos_cis_img: torch.Tensor
+        sin_cis_img: torch.Tensor
+        return_dict: bool
+            Whether to return a dictionary.
+        """
+        condition = hint
+        if condition.shape[0] == 1:
+            condition = torch.repeat_interleave(condition, x.shape[0], dim=0)
+
+        text_states = context  # 2,77,1024
+        text_states_t5 = encoder_hidden_states_t5  # 2,256,2048
+        text_states_mask = text_embedding_mask.bool()  # 2,77
+        text_states_t5_mask = text_embedding_mask_t5.bool()  # 2,256
+        b_t5, l_t5, c_t5 = text_states_t5.shape
+        text_states_t5 = self.mlp_t5(text_states_t5.view(-1, c_t5)).view(b_t5, l_t5, -1)
+
+        padding = comfy.ops.cast_to_input(self.text_embedding_padding, text_states)
+
+        text_states[:, -self.text_len :] = torch.where(
+            text_states_mask[:, -self.text_len :].unsqueeze(2),
+            text_states[:, -self.text_len :],
+            padding[: self.text_len],
+        )
+        text_states_t5[:, -self.text_len_t5 :] = torch.where(
+            text_states_t5_mask[:, -self.text_len_t5 :].unsqueeze(2),
+            text_states_t5[:, -self.text_len_t5 :],
+            padding[self.text_len :],
+        )
+
+        text_states = torch.cat([text_states, text_states_t5], dim=1)  # 2,205，1024
+
+        # _, _, oh, ow = x.shape
+        # th, tw = oh // self.patch_size, ow // self.patch_size
+
+        # Get image RoPE embedding according to `reso`lution.
+        freqs_cis_img = calc_rope(
+            x, self.patch_size, self.hidden_size // self.num_heads
+        )  # (cos_cis_img, sin_cis_img)
+
+        # ========================= Build time and image embedding =========================
+        t = self.t_embedder(timesteps, dtype=self.dtype)
+        x = self.x_embedder(x)
+
+        # ========================= Concatenate all extra vectors =========================
+        # Build text tokens with pooling
+        extra_vec = self.pooler(encoder_hidden_states_t5)
+
+        # Build image meta size tokens if applicable
+        # if image_meta_size is not None:
+        #     image_meta_size = timestep_embedding(image_meta_size.view(-1), 256)   # [B * 6, 256]
+        #     if image_meta_size.dtype != self.dtype:
+        #         image_meta_size = image_meta_size.half()
+        #     image_meta_size = image_meta_size.view(-1, 6 * 256)
+        #     extra_vec = torch.cat([extra_vec, image_meta_size], dim=1)  # [B, D + 6 * 256]
+
+        # Build style tokens
+        if style is not None:
+            style_embedding = self.style_embedder(style)
+            extra_vec = torch.cat([extra_vec, style_embedding], dim=1)
+
+        # Concatenate all extra vectors
+        c = t + self.extra_embedder(extra_vec)  # [B, D]
+
+        # ========================= Deal with Condition =========================
+        condition = self.x_embedder(condition)
+
+        # ========================= Forward pass through HunYuanDiT blocks =========================
+        controls = []
+        x = x + self.before_proj(condition)  # add condition
+        for layer, block in enumerate(self.blocks):
+            x = block(x, c, text_states, freqs_cis_img)
+            controls.append(self.after_proj_list[layer](x))  # zero linear for output
+
+        return {"output": controls}

comfy/ldm/hydit/models.py

@@ -21,6 +21,7 @@ def calc_rope(x, patch_size, head_size):
     sub_args = [start, stop, (th, tw)]
     # head_size = HUNYUAN_DIT_CONFIG['DiT-g/2']['hidden_size'] // HUNYUAN_DIT_CONFIG['DiT-g/2']['num_heads']
     rope = get_2d_rotary_pos_embed(head_size, *sub_args)
+    rope = (rope[0].to(x), rope[1].to(x))
     return rope
 
 
@@ -91,6 +92,8 @@ class HunYuanDiTBlock(nn.Module):
         # Long Skip Connection
         if self.skip_linear is not None:
             cat = torch.cat([x, skip], dim=-1)
+            if cat.dtype != x.dtype:
+                cat = cat.to(x.dtype)
             cat = self.skip_norm(cat)
             x = self.skip_linear(cat)
 
@@ -362,6 +365,8 @@ class HunYuanDiT(nn.Module):
         c = t + self.extra_embedder(extra_vec)  # [B, D]
 
         controls = None
+        if control:
+            controls = control.get("output", None)
         # ========================= Forward pass through HunYuanDiT blocks =========================
         skips = []
         for layer, block in enumerate(self.blocks):

comfy/ldm/modules/attention.py

@@ -358,7 +358,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
             disabled_xformers = True
 
     if disabled_xformers:
-        return attention_pytorch(q, k, v, heads, mask)
+        return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape)
 
     if skip_reshape:
          q, k, v = map(

comfy/lora.py

@@ -1,3 +1,21 @@
+"""
+    This file is part of ComfyUI.
+    Copyright (C) 2024 Comfy
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
 import comfy.utils
 import logging
 
@@ -218,11 +236,17 @@ def model_lora_keys_clip(model, key_map={}):
                     lora_key = "lora_prior_te_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #cascade lora: TODO put lora key prefix in the model config
                     key_map[lora_key] = k
 
-    for k in sdk: #OneTrainer SD3 lora
-        if k.startswith("t5xxl.transformer.") and k.endswith(".weight"):
-            l_key = k[len("t5xxl.transformer."):-len(".weight")]
-            lora_key = "lora_te3_{}".format(l_key.replace(".", "_"))
-            key_map[lora_key] = k
+    for k in sdk:
+        if k.endswith(".weight"):
+            if k.startswith("t5xxl.transformer."):#OneTrainer SD3 lora
+                l_key = k[len("t5xxl.transformer."):-len(".weight")]
+                lora_key = "lora_te3_{}".format(l_key.replace(".", "_"))
+                key_map[lora_key] = k
+            elif k.startswith("hydit_clip.transformer.bert."): #HunyuanDiT Lora
+                l_key = k[len("hydit_clip.transformer.bert."):-len(".weight")]
+                lora_key = "lora_te1_{}".format(l_key.replace(".", "_"))
+                key_map[lora_key] = k
+
 
     k = "clip_g.transformer.text_projection.weight"
     if k in sdk:
@@ -245,6 +269,7 @@ def model_lora_keys_unet(model, key_map={}):
             key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")
             key_map["lora_unet_{}".format(key_lora)] = k
             key_map["lora_prior_unet_{}".format(key_lora)] = k #cascade lora: TODO put lora key prefix in the model config
+            key_map["{}".format(k[:-len(".weight")])] = k #generic lora format without any weird key names
 
     diffusers_keys = comfy.utils.unet_to_diffusers(model.model_config.unet_config)
     for k in diffusers_keys:

comfy/model_base.py

@@ -1,3 +1,21 @@
+"""
+    This file is part of ComfyUI.
+    Copyright (C) 2024 Comfy
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
 import torch
 import logging
 from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
@@ -77,10 +95,13 @@ class BaseModel(torch.nn.Module):
         self.device = device
 
         if not unet_config.get("disable_unet_model_creation", False):
-            if self.manual_cast_dtype is not None:
-                operations = comfy.ops.manual_cast
+            if model_config.custom_operations is None:
+                if self.manual_cast_dtype is not None:
+                    operations = comfy.ops.manual_cast
+                else:
+                    operations = comfy.ops.disable_weight_init
             else:
-                operations = comfy.ops.disable_weight_init
+                operations = model_config.custom_operations
             self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
             if comfy.model_management.force_channels_last():
                 self.diffusion_model.to(memory_format=torch.channels_last)

comfy/model_detection.py

@@ -137,8 +137,8 @@ def detect_unet_config(state_dict, key_prefix):
         dit_config["hidden_size"] = 3072
         dit_config["mlp_ratio"] = 4.0
         dit_config["num_heads"] = 24
-        dit_config["depth"] = 19
-        dit_config["depth_single_blocks"] = 38
+        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
@@ -495,7 +495,12 @@ def model_config_from_diffusers_unet(state_dict):
 def convert_diffusers_mmdit(state_dict, output_prefix=""):
     out_sd = {}
 
-    if 'transformer_blocks.0.attn.add_q_proj.weight' in state_dict: #SD3
+    if 'transformer_blocks.0.attn.norm_added_k.weight' in state_dict: #Flux
+        depth = count_blocks(state_dict, 'transformer_blocks.{}.')
+        depth_single_blocks = count_blocks(state_dict, 'single_transformer_blocks.{}.')
+        hidden_size = state_dict["x_embedder.bias"].shape[0]
+        sd_map = comfy.utils.flux_to_diffusers({"depth": depth, "depth_single_blocks": depth_single_blocks, "hidden_size": hidden_size}, output_prefix=output_prefix)
+    elif 'transformer_blocks.0.attn.add_q_proj.weight' in state_dict: #SD3
         num_blocks = count_blocks(state_dict, 'transformer_blocks.{}.')
         depth = state_dict["pos_embed.proj.weight"].shape[0] // 64
         sd_map = comfy.utils.mmdit_to_diffusers({"depth": depth, "num_blocks": num_blocks}, output_prefix=output_prefix)
@@ -521,7 +526,12 @@ def convert_diffusers_mmdit(state_dict, output_prefix=""):
                     old_weight = out_sd.get(t[0], None)
                     if old_weight is None:
                         old_weight = torch.empty_like(weight)
-                        old_weight = old_weight.repeat([3] + [1] * (len(old_weight.shape) - 1))
+                    if old_weight.shape[offset[0]] < offset[1] + offset[2]:
+                        exp = list(weight.shape)
+                        exp[offset[0]] = offset[1] + offset[2]
+                        new = torch.empty(exp, device=weight.device, dtype=weight.dtype)
+                        new[:old_weight.shape[0]] = old_weight
+                        old_weight = new
 
                     w = old_weight.narrow(offset[0], offset[1], offset[2])
                 else:

comfy/model_management.py

@@ -1,3 +1,21 @@
+"""
+    This file is part of ComfyUI.
+    Copyright (C) 2024 Comfy
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
 import psutil
 import logging
 from enum import Enum
@@ -273,9 +291,12 @@ class LoadedModel:
     def model_memory(self):
         return self.model.model_size()
 
+    def model_offloaded_memory(self):
+        return self.model.model_size() - self.model.loaded_size()
+
     def model_memory_required(self, device):
         if device == self.model.current_loaded_device():
-            return 0
+            return self.model_offloaded_memory()
         else:
             return self.model_memory()
 
@@ -287,15 +308,21 @@ class LoadedModel:
 
         load_weights = not self.weights_loaded
 
-        try:
-            if lowvram_model_memory > 0 and load_weights:
-                self.real_model = self.model.patch_model_lowvram(device_to=patch_model_to, lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights)
-            else:
-                self.real_model = self.model.patch_model(device_to=patch_model_to, patch_weights=load_weights)
-        except Exception as e:
-            self.model.unpatch_model(self.model.offload_device)
-            self.model_unload()
-            raise e
+        if self.model.loaded_size() > 0:
+            use_more_vram = lowvram_model_memory
+            if use_more_vram == 0:
+                use_more_vram = 1e32
+            self.model_use_more_vram(use_more_vram)
+        else:
+            try:
+                if lowvram_model_memory > 0 and load_weights:
+                    self.real_model = self.model.patch_model_lowvram(device_to=patch_model_to, lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights)
+                else:
+                    self.real_model = self.model.patch_model(device_to=patch_model_to, patch_weights=load_weights)
+            except Exception as e:
+                self.model.unpatch_model(self.model.offload_device)
+                self.model_unload()
+                raise e
 
         if is_intel_xpu() and not args.disable_ipex_optimize:
             self.real_model = ipex.optimize(self.real_model.eval(), graph_mode=True, concat_linear=True)
@@ -304,19 +331,42 @@ class LoadedModel:
         return self.real_model
 
     def should_reload_model(self, force_patch_weights=False):
-        if force_patch_weights and self.model.lowvram_patch_counter > 0:
+        if force_patch_weights and self.model.lowvram_patch_counter() > 0:
             return True
         return False
 
-    def model_unload(self, unpatch_weights=True):
+    def model_unload(self, memory_to_free=None, unpatch_weights=True):
+        if memory_to_free is not None:
+            if memory_to_free < self.model.loaded_size():
+                freed = self.model.partially_unload(self.model.offload_device, memory_to_free)
+                if freed >= memory_to_free:
+                    return False
         self.model.unpatch_model(self.model.offload_device, unpatch_weights=unpatch_weights)
         self.model.model_patches_to(self.model.offload_device)
         self.weights_loaded = self.weights_loaded and not unpatch_weights
         self.real_model = None
+        return True
+
+    def model_use_more_vram(self, extra_memory):
+        return self.model.partially_load(self.device, extra_memory)
 
     def __eq__(self, other):
         return self.model is other.model
 
+def use_more_memory(extra_memory, loaded_models, device):
+    for m in loaded_models:
+        if m.device == device:
+            extra_memory -= m.model_use_more_vram(extra_memory)
+            if extra_memory <= 0:
+                break
+
+def offloaded_memory(loaded_models, device):
+    offloaded_mem = 0
+    for m in loaded_models:
+        if m.device == device:
+            offloaded_mem += m.model_offloaded_memory()
+    return offloaded_mem
+
 def minimum_inference_memory():
     return (1024 * 1024 * 1024) * 1.2
 
@@ -363,11 +413,15 @@ def free_memory(memory_required, device, keep_loaded=[]):
 
     for x in sorted(can_unload):
         i = x[-1]
+        memory_to_free = None
         if not DISABLE_SMART_MEMORY:
-            if get_free_memory(device) > memory_required:
+            free_mem = get_free_memory(device)
+            if free_mem > memory_required:
                 break
-        current_loaded_models[i].model_unload()
-        unloaded_model.append(i)
+            memory_to_free = memory_required - free_mem
+        logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}")
+        if current_loaded_models[i].model_unload(memory_to_free):
+            unloaded_model.append(i)
 
     for i in sorted(unloaded_model, reverse=True):
         unloaded_models.append(current_loaded_models.pop(i))
@@ -381,15 +435,15 @@ def free_memory(memory_required, device, keep_loaded=[]):
                 soft_empty_cache()
     return unloaded_models
 
-def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None):
+def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False):
     global vram_state
 
     inference_memory = minimum_inference_memory()
-    extra_mem = max(inference_memory, memory_required)
+    extra_mem = max(inference_memory, memory_required + 300 * 1024 * 1024)
     if minimum_memory_required is None:
         minimum_memory_required = extra_mem
     else:
-        minimum_memory_required = max(inference_memory, minimum_memory_required)
+        minimum_memory_required = max(inference_memory, minimum_memory_required + 300 * 1024 * 1024)
 
     models = set(models)
 
@@ -422,12 +476,14 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
         devs = set(map(lambda a: a.device, models_already_loaded))
         for d in devs:
             if d != torch.device("cpu"):
-                free_memory(extra_mem, d, models_already_loaded)
+                free_memory(extra_mem + offloaded_memory(models_already_loaded, d), d, models_already_loaded)
                 free_mem = get_free_memory(d)
                 if free_mem < minimum_memory_required:
                     logging.info("Unloading models for lowram load.") #TODO: partial model unloading when this case happens, also handle the opposite case where models can be unlowvramed.
                     models_to_load = free_memory(minimum_memory_required, d)
                     logging.info("{} models unloaded.".format(len(models_to_load)))
+                else:
+                    use_more_memory(free_mem - minimum_memory_required, models_already_loaded, d)
         if len(models_to_load) == 0:
             return
 
@@ -435,18 +491,21 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
 
     total_memory_required = {}
     for loaded_model in models_to_load:
-        if unload_model_clones(loaded_model.model, unload_weights_only=True, force_unload=False) == True:#unload clones where the weights are different
-            total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
+        unload_model_clones(loaded_model.model, unload_weights_only=True, force_unload=False) #unload clones where the weights are different
+        total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
 
-    for device in total_memory_required:
-        if device != torch.device("cpu"):
-            free_memory(total_memory_required[device] * 1.3 + extra_mem, device, models_already_loaded)
+    for loaded_model in models_already_loaded:
+        total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
 
     for loaded_model in models_to_load:
         weights_unloaded = unload_model_clones(loaded_model.model, unload_weights_only=False, force_unload=False) #unload the rest of the clones where the weights can stay loaded
         if weights_unloaded is not None:
             loaded_model.weights_loaded = not weights_unloaded
 
+    for device in total_memory_required:
+        if device != torch.device("cpu"):
+            free_memory(total_memory_required[device] * 1.1 + extra_mem, device, models_already_loaded)
+
     for loaded_model in models_to_load:
         model = loaded_model.model
         torch_dev = model.load_device
@@ -455,7 +514,7 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
         else:
             vram_set_state = vram_state
         lowvram_model_memory = 0
-        if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM):
+        if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM) and not force_full_load:
             model_size = loaded_model.model_memory_required(torch_dev)
             current_free_mem = get_free_memory(torch_dev)
             lowvram_model_memory = max(64 * (1024 * 1024), (current_free_mem - minimum_memory_required), min(current_free_mem * 0.4, current_free_mem - minimum_inference_memory()))
@@ -467,6 +526,14 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
 
         cur_loaded_model = loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
         current_loaded_models.insert(0, loaded_model)
+
+
+    devs = set(map(lambda a: a.device, models_already_loaded))
+    for d in devs:
+        if d != torch.device("cpu"):
+            free_mem = get_free_memory(d)
+            if free_mem > minimum_memory_required:
+                use_more_memory(free_mem - minimum_memory_required, models_already_loaded, d)
     return
 
 
@@ -562,12 +629,22 @@ def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, tor
         if model_params * 2 > free_model_memory:
             return fp8_dtype
 
-    if should_use_fp16(device=device, model_params=model_params, manual_cast=True):
-        if torch.float16 in supported_dtypes:
-            return torch.float16
-    if should_use_bf16(device, model_params=model_params, manual_cast=True):
-        if torch.bfloat16 in supported_dtypes:
-            return torch.bfloat16
+    for dt in supported_dtypes:
+        if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params):
+            if torch.float16 in supported_dtypes:
+                return torch.float16
+        if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params):
+            if torch.bfloat16 in supported_dtypes:
+                return torch.bfloat16
+
+    for dt in supported_dtypes:
+        if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params, manual_cast=True):
+            if torch.float16 in supported_dtypes:
+                return torch.float16
+        if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params, manual_cast=True):
+            if torch.bfloat16 in supported_dtypes:
+                return torch.bfloat16
+
     return torch.float32
 
 # None means no manual cast
@@ -583,13 +660,13 @@ def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.flo
     if bf16_supported and weight_dtype == torch.bfloat16:
         return None
 
-    if fp16_supported and torch.float16 in supported_dtypes:
-        return torch.float16
+    for dt in supported_dtypes:
+        if dt == torch.float16 and fp16_supported:
+            return torch.float16
+        if dt == torch.bfloat16 and bf16_supported:
+            return torch.bfloat16
 
-    elif bf16_supported and torch.bfloat16 in supported_dtypes:
-        return torch.bfloat16
-    else:
-        return torch.float32
+    return torch.float32
 
 def text_encoder_offload_device():
     if args.gpu_only:
@@ -608,6 +685,20 @@ def text_encoder_device():
     else:
         return torch.device("cpu")
 
+def text_encoder_initial_device(load_device, offload_device, model_size=0):
+    if load_device == offload_device or model_size <= 1024 * 1024 * 1024:
+        return offload_device
+
+    if is_device_mps(load_device):
+        return offload_device
+
+    mem_l = get_free_memory(load_device)
+    mem_o = get_free_memory(offload_device)
+    if mem_l > (mem_o * 0.5) and model_size * 1.2 < mem_l:
+        return load_device
+    else:
+        return offload_device
+
 def text_encoder_dtype(device=None):
     if args.fp8_e4m3fn_text_enc:
         return torch.float8_e4m3fn

comfy/model_patcher.py

@@ -1,8 +1,27 @@
+"""
+    This file is part of ComfyUI.
+    Copyright (C) 2024 Comfy
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
 import torch
 import copy
 import inspect
 import logging
 import uuid
+import collections
 
 import comfy.utils
 import comfy.model_management
@@ -63,12 +82,27 @@ def set_model_options_pre_cfg_function(model_options, pre_cfg_function, disable_
         model_options["disable_cfg1_optimization"] = True
     return model_options
 
+def wipe_lowvram_weight(m):
+    if hasattr(m, "prev_comfy_cast_weights"):
+        m.comfy_cast_weights = m.prev_comfy_cast_weights
+        del m.prev_comfy_cast_weights
+    m.weight_function = None
+    m.bias_function = None
+
+class LowVramPatch:
+    def __init__(self, key, model_patcher):
+        self.key = key
+        self.model_patcher = model_patcher
+    def __call__(self, weight):
+        return self.model_patcher.calculate_weight(self.model_patcher.patches[self.key], weight, self.key)
+
+
 class ModelPatcher:
     def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False):
         self.size = size
         self.model = model
         if not hasattr(self.model, 'device'):
-            logging.info("Model doesn't have a device attribute.")
+            logging.debug("Model doesn't have a device attribute.")
             self.model.device = offload_device
         elif self.model.device is None:
             self.model.device = offload_device
@@ -82,16 +116,29 @@ class ModelPatcher:
         self.load_device = load_device
         self.offload_device = offload_device
         self.weight_inplace_update = weight_inplace_update
-        self.model_lowvram = False
-        self.lowvram_patch_counter = 0
         self.patches_uuid = uuid.uuid4()
 
+        if not hasattr(self.model, 'model_loaded_weight_memory'):
+            self.model.model_loaded_weight_memory = 0
+
+        if not hasattr(self.model, 'lowvram_patch_counter'):
+            self.model.lowvram_patch_counter = 0
+
+        if not hasattr(self.model, 'model_lowvram'):
+            self.model.model_lowvram = False
+
     def model_size(self):
         if self.size > 0:
             return self.size
         self.size = comfy.model_management.module_size(self.model)
         return self.size
 
+    def loaded_size(self):
+        return self.model.model_loaded_weight_memory
+
+    def lowvram_patch_counter(self):
+        return self.model.lowvram_patch_counter
+
     def clone(self):
         n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, weight_inplace_update=self.weight_inplace_update)
         n.patches = {}
@@ -265,16 +312,16 @@ class ModelPatcher:
                     sd.pop(k)
         return sd
 
-    def patch_weight_to_device(self, key, device_to=None):
+    def patch_weight_to_device(self, key, device_to=None, inplace_update=False):
         if key not in self.patches:
             return
 
         weight = comfy.utils.get_attr(self.model, key)
 
-        inplace_update = self.weight_inplace_update
+        inplace_update = self.weight_inplace_update or inplace_update
 
         if key not in self.backup:
-            self.backup[key] = weight.to(device=self.offload_device, copy=inplace_update)
+            self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update)
 
         if device_to is not None:
             temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
@@ -304,28 +351,24 @@ class ModelPatcher:
             if device_to is not None:
                 self.model.to(device_to)
                 self.model.device = device_to
+                self.model.model_loaded_weight_memory = self.model_size()
 
         return self.model
 
-    def patch_model_lowvram(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False):
-        self.patch_model(device_to, patch_weights=False)
-
-        logging.info("loading in lowvram mode {}".format(lowvram_model_memory/(1024 * 1024)))
-        class LowVramPatch:
-            def __init__(self, key, model_patcher):
-                self.key = key
-                self.model_patcher = model_patcher
-            def __call__(self, weight):
-                return self.model_patcher.calculate_weight(self.model_patcher.patches[self.key], weight, self.key)
-
+    def lowvram_load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False):
         mem_counter = 0
         patch_counter = 0
+        lowvram_counter = 0
         for n, m in self.model.named_modules():
             lowvram_weight = False
-            if hasattr(m, "comfy_cast_weights"):
+
+            if not full_load and hasattr(m, "comfy_cast_weights"):
                 module_mem = comfy.model_management.module_size(m)
                 if mem_counter + module_mem >= lowvram_model_memory:
                     lowvram_weight = True
+                    lowvram_counter += 1
+                    if m.comfy_cast_weights:
+                        continue
 
             weight_key = "{}.weight".format(n)
             bias_key = "{}.bias".format(n)
@@ -347,16 +390,40 @@ class ModelPatcher:
                 m.prev_comfy_cast_weights = m.comfy_cast_weights
                 m.comfy_cast_weights = True
             else:
+                if hasattr(m, "comfy_cast_weights"):
+                    if m.comfy_cast_weights:
+                        wipe_lowvram_weight(m)
+
                 if hasattr(m, "weight"):
-                    self.patch_weight_to_device(weight_key, device_to)
-                    self.patch_weight_to_device(bias_key, device_to)
-                    m.to(device_to)
                     mem_counter += comfy.model_management.module_size(m)
+                    param = list(m.parameters())
+                    if len(param) > 0:
+                        weight = param[0]
+                        if weight.device == device_to:
+                            continue
+
+                    weight_to = None
+                    if full_load:#TODO
+                        weight_to = device_to
+                    self.patch_weight_to_device(weight_key, device_to=weight_to) #TODO: speed this up without OOM
+                    self.patch_weight_to_device(bias_key, device_to=weight_to)
+                    m.to(device_to)
                     logging.debug("lowvram: loaded module regularly {} {}".format(n, m))
 
-        self.model_lowvram = True
-        self.lowvram_patch_counter = patch_counter
+        if lowvram_counter > 0:
+            logging.info("loaded partially {} {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), patch_counter))
+            self.model.model_lowvram = True
+        else:
+            logging.info("loaded completely {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024)))
+            self.model.model_lowvram = False
+        self.model.lowvram_patch_counter += patch_counter
         self.model.device = device_to
+        self.model.model_loaded_weight_memory = mem_counter
+
+
+    def patch_model_lowvram(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False):
+        self.patch_model(device_to, patch_weights=False)
+        self.lowvram_load(device_to, lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights)
         return self.model
 
     def calculate_weight(self, patches, weight, key):
@@ -529,31 +596,28 @@ class ModelPatcher:
 
     def unpatch_model(self, device_to=None, unpatch_weights=True):
         if unpatch_weights:
-            if self.model_lowvram:
+            if self.model.model_lowvram:
                 for m in self.model.modules():
-                    if hasattr(m, "prev_comfy_cast_weights"):
-                        m.comfy_cast_weights = m.prev_comfy_cast_weights
-                        del m.prev_comfy_cast_weights
-                    m.weight_function = None
-                    m.bias_function = None
+                    wipe_lowvram_weight(m)
 
-                self.model_lowvram = False
-                self.lowvram_patch_counter = 0
+                self.model.model_lowvram = False
+                self.model.lowvram_patch_counter = 0
 
             keys = list(self.backup.keys())
 
-            if self.weight_inplace_update:
-                for k in keys:
-                    comfy.utils.copy_to_param(self.model, k, self.backup[k])
-            else:
-                for k in keys:
-                    comfy.utils.set_attr_param(self.model, k, self.backup[k])
+            for k in keys:
+                bk = self.backup[k]
+                if bk.inplace_update:
+                    comfy.utils.copy_to_param(self.model, k, bk.weight)
+                else:
+                    comfy.utils.set_attr_param(self.model, k, bk.weight)
 
             self.backup.clear()
 
             if device_to is not None:
                 self.model.to(device_to)
                 self.model.device = device_to
+            self.model.model_loaded_weight_memory = 0
 
         keys = list(self.object_patches_backup.keys())
         for k in keys:
@@ -561,5 +625,60 @@ class ModelPatcher:
 
         self.object_patches_backup.clear()
 
+    def partially_unload(self, device_to, memory_to_free=0):
+        memory_freed = 0
+        patch_counter = 0
+
+        for n, m in list(self.model.named_modules())[::-1]:
+            if memory_to_free < memory_freed:
+                break
+
+            shift_lowvram = False
+            if hasattr(m, "comfy_cast_weights"):
+                module_mem = comfy.model_management.module_size(m)
+                weight_key = "{}.weight".format(n)
+                bias_key = "{}.bias".format(n)
+
+
+                if m.weight is not None and m.weight.device != device_to:
+                    for key in [weight_key, bias_key]:
+                        bk = self.backup.get(key, None)
+                        if bk is not None:
+                            if bk.inplace_update:
+                                comfy.utils.copy_to_param(self.model, key, bk.weight)
+                            else:
+                                comfy.utils.set_attr_param(self.model, key, bk.weight)
+                            self.backup.pop(key)
+
+                    m.to(device_to)
+                    if weight_key in self.patches:
+                        m.weight_function = LowVramPatch(weight_key, self)
+                        patch_counter += 1
+                    if bias_key in self.patches:
+                        m.bias_function = LowVramPatch(bias_key, self)
+                        patch_counter += 1
+
+                    m.prev_comfy_cast_weights = m.comfy_cast_weights
+                    m.comfy_cast_weights = True
+                    memory_freed += module_mem
+                    logging.debug("freed {}".format(n))
+
+        self.model.model_lowvram = True
+        self.model.lowvram_patch_counter += patch_counter
+        self.model.model_loaded_weight_memory -= memory_freed
+        return memory_freed
+
+    def partially_load(self, device_to, extra_memory=0):
+        self.unpatch_model(unpatch_weights=False)
+        self.patch_model(patch_weights=False)
+        full_load = False
+        if self.model.model_lowvram == False:
+            return 0
+        if self.model.model_loaded_weight_memory + extra_memory > self.model_size():
+            full_load = True
+        current_used = self.model.model_loaded_weight_memory
+        self.lowvram_load(device_to, lowvram_model_memory=current_used + extra_memory, full_load=full_load)
+        return self.model.model_loaded_weight_memory - current_used
+
     def current_loaded_device(self):
         return self.model.device

comfy/sd.py

@@ -62,7 +62,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={}):
+    def __init__(self, target=None, embedding_directory=None, no_init=False, tokenizer_data={}, parameters=0):
         if no_init:
             return
         params = target.params.copy()
@@ -71,20 +71,24 @@ class CLIP:
 
         load_device = model_management.text_encoder_device()
         offload_device = model_management.text_encoder_offload_device()
-        params['device'] = offload_device
         dtype = model_management.text_encoder_dtype(load_device)
         params['dtype'] = dtype
-
+        params['device'] = model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype))
         self.cond_stage_model = clip(**(params))
 
         for dt in self.cond_stage_model.dtypes:
             if not model_management.supports_cast(load_device, dt):
                 load_device = offload_device
+                if params['device'] != offload_device:
+                    self.cond_stage_model.to(offload_device)
+                    logging.warning("Had to shift TE back.")
 
         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)
+        if params['device'] == load_device:
+            model_management.load_models_gpu([self.patcher], force_full_load=True)
         self.layer_idx = None
-        logging.debug("CLIP model load device: {}, offload device: {}".format(load_device, offload_device))
+        logging.debug("CLIP model load device: {}, offload device: {}, current: {}".format(load_device, offload_device, params['device']))
 
     def clone(self):
         n = CLIP(no_init=True)
@@ -456,7 +460,11 @@ def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DI
         clip_target.clip = comfy.text_encoders.sd3_clip.SD3ClipModel
         clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
 
-    clip = CLIP(clip_target, embedding_directory=embedding_directory)
+    parameters = 0
+    for c in clip_data:
+        parameters += comfy.utils.calculate_parameters(c)
+
+    clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters)
     for c in clip_data:
         m, u = clip.load_sd(c)
         if len(m) > 0:
@@ -498,15 +506,19 @@ def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_cl
 
     return (model, clip, vae)
 
-def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True):
+def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True, model_options={}):
     sd = comfy.utils.load_torch_file(ckpt_path)
-    sd_keys = sd.keys()
+    out = load_state_dict_guess_config(sd, output_vae, output_clip, output_clipvision, embedding_directory, output_model, model_options)
+    if out is None:
+        raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
+    return out
+
+def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True, model_options={}):
     clip = None
     clipvision = None
     vae = None
     model = None
     model_patcher = None
-    clip_target = None
 
     diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
     parameters = comfy.utils.calculate_parameters(sd, diffusion_model_prefix)
@@ -515,13 +527,18 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
 
     model_config = model_detection.model_config_from_unet(sd, diffusion_model_prefix)
     if model_config is None:
-        raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
+        return None
 
     unet_weight_dtype = list(model_config.supported_inference_dtypes)
     if weight_dtype is not None:
         unet_weight_dtype.append(weight_dtype)
 
-    unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype)
+    model_config.custom_operations = model_options.get("custom_operations", None)
+    unet_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)
+
     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)
 
@@ -545,7 +562,8 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
         if clip_target is not None:
             clip_sd = model_config.process_clip_state_dict(sd)
             if len(clip_sd) > 0:
-                clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd)
+                parameters = comfy.utils.calculate_parameters(clip_sd)
+                clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd, parameters=parameters)
                 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))
@@ -567,12 +585,13 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
         model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device())
         if inital_load_device != torch.device("cpu"):
             logging.info("loaded straight to GPU")
-            model_management.load_model_gpu(model_patcher)
+            model_management.load_models_gpu([model_patcher], force_full_load=True)
 
     return (model_patcher, clip, vae, clipvision)
 
 
-def load_unet_state_dict(sd, dtype=None): #load unet in diffusers or regular format
+def load_diffusion_model_state_dict(sd, model_options={}): #load unet in diffusers or regular format
+    dtype = model_options.get("dtype", None)
 
     #Allow loading unets from checkpoint files
     diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
@@ -614,6 +633,7 @@ def load_unet_state_dict(sd, dtype=None): #load unet in diffusers or regular for
 
     manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
     model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+    model_config.custom_operations = model_options.get("custom_operations", None)
     model = model_config.get_model(new_sd, "")
     model = model.to(offload_device)
     model.load_model_weights(new_sd, "")
@@ -622,14 +642,23 @@ def load_unet_state_dict(sd, dtype=None): #load unet in diffusers or regular for
         logging.info("left over keys in unet: {}".format(left_over))
     return comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device)
 
-def load_unet(unet_path, dtype=None):
+
+def load_diffusion_model(unet_path, model_options={}):
     sd = comfy.utils.load_torch_file(unet_path)
-    model = load_unet_state_dict(sd, dtype=dtype)
+    model = load_diffusion_model_state_dict(sd, model_options=model_options)
     if model is None:
         logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path))
         raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
     return model
 
+def load_unet(unet_path, dtype=None):
+    print("WARNING: the load_unet function has been deprecated and will be removed please switch to: load_diffusion_model")
+    return load_diffusion_model(unet_path, model_options={"dtype": dtype})
+
+def load_unet_state_dict(sd, dtype=None):
+    print("WARNING: the load_unet_state_dict function has been deprecated and will be removed please switch to: load_diffusion_model_state_dict")
+    return load_diffusion_model_state_dict(sd, model_options={"dtype": dtype})
+
 def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None, extra_keys={}):
     clip_sd = None
     load_models = [model]

comfy/sd1_clip.py

@@ -313,17 +313,14 @@ def expand_directory_list(directories):
             dirs.add(root)
     return list(dirs)
 
-def bundled_embed(embed, key): #bundled embedding in lora format
+def bundled_embed(embed, prefix, suffix): #bundled embedding in lora format
     i = 0
     out_list = []
-    while True:
-        i += 1
-        k = key.format(i)
-        w = embed.get(k, None)
-        if w is None:
-            break
-        else:
-            out_list.append(w)
+    for k in embed:
+        if k.startswith(prefix) and k.endswith(suffix):
+            out_list.append(embed[k])
+    if len(out_list) == 0:
+        return None
 
     return torch.cat(out_list, dim=0)
 
@@ -392,13 +389,13 @@ def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=No
             embed_out = torch.cat(out_list, dim=0)
         elif embed_key is not None and embed_key in embed:
             embed_out = embed[embed_key]
-        elif 'bundle_emb.place1.string_to_param.*' in embed:
-            embed_out = bundled_embed(embed, 'bundle_emb.place{}.string_to_param.*')
-        elif 'bundle_emb.place1.{}'.format(embed_key) in embed:
-            embed_out = bundled_embed(embed, 'bundle_emb.place{}.{}'.format('{}', embed_key))
         else:
-            values = embed.values()
-            embed_out = next(iter(values))
+            embed_out = bundled_embed(embed, 'bundle_emb.', '.string_to_param.*')
+            if embed_out is None:
+                embed_out = bundled_embed(embed, 'bundle_emb.', '.{}'.format(embed_key))
+            if embed_out is None:
+                values = embed.values()
+                embed_out = next(iter(values))
     return embed_out
 
 class SDTokenizer:

comfy/supported_models.py

@@ -181,7 +181,7 @@ class SDXL(supported_models_base.BASE):
 
     latent_format = latent_formats.SDXL
 
-    memory_usage_factor = 0.7
+    memory_usage_factor = 0.8
 
     def model_type(self, state_dict, prefix=""):
         if 'edm_mean' in state_dict and 'edm_std' in state_dict: #Playground V2.5
@@ -642,7 +642,7 @@ class Flux(supported_models_base.BASE):
 
     memory_usage_factor = 2.8
 
-    supported_inference_dtypes = [torch.bfloat16, torch.float32]
+    supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
 
     vae_key_prefix = ["vae."]
     text_encoder_key_prefix = ["text_encoders."]

comfy/supported_models_base.py

@@ -1,3 +1,21 @@
+"""
+    This file is part of ComfyUI.
+    Copyright (C) 2024 Comfy
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
 import torch
 from . import model_base
 from . import utils
@@ -30,6 +48,7 @@ class BASE:
     memory_usage_factor = 2.0
 
     manual_cast_dtype = None
+    custom_operations = None
 
     @classmethod
     def matches(s, unet_config, state_dict=None):

comfy/utils.py

@@ -1,3 +1,22 @@
+"""
+    This file is part of ComfyUI.
+    Copyright (C) 2024 Comfy
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+
 import torch
 import math
 import struct
@@ -432,8 +451,33 @@ def flux_to_diffusers(mmdit_config, output_prefix=""):
             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 = {"attn.to_out.0.weight": "img_attn.proj.weight",
-                     "attn.to_out.0.bias": "img_attn.proj.bias",
+            k = "{}.attn.".format(prefix_from)
+            qkv = "{}.txt_attn.qkv.{}".format(prefix_to, end)
+            key_map["{}add_q_proj.{}".format(k, end)] = (qkv, (0, 0, hidden_size))
+            key_map["{}add_k_proj.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))
+            key_map["{}add_v_proj.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))
+
+        block_map = {
+                        "attn.to_out.0.weight": "img_attn.proj.weight",
+                        "attn.to_out.0.bias": "img_attn.proj.bias",
+                        "norm1.linear.weight": "img_mod.lin.weight",
+                        "norm1.linear.bias": "img_mod.lin.bias",
+                        "norm1_context.linear.weight": "txt_mod.lin.weight",
+                        "norm1_context.linear.bias": "txt_mod.lin.bias",
+                        "attn.to_add_out.weight": "txt_attn.proj.weight",
+                        "attn.to_add_out.bias": "txt_attn.proj.bias",
+                        "ff.net.0.proj.weight": "img_mlp.0.weight",
+                        "ff.net.0.proj.bias": "img_mlp.0.bias",
+                        "ff.net.2.weight": "img_mlp.2.weight",
+                        "ff.net.2.bias": "img_mlp.2.bias",
+                        "ff_context.net.0.proj.weight": "txt_mlp.0.weight",
+                        "ff_context.net.0.proj.bias": "txt_mlp.0.bias",
+                        "ff_context.net.2.weight": "txt_mlp.2.weight",
+                        "ff_context.net.2.bias": "txt_mlp.2.bias",
+                        "attn.norm_q.weight": "img_attn.norm.query_norm.scale",
+                        "attn.norm_k.weight": "img_attn.norm.key_norm.scale",
+                        "attn.norm_added_q.weight": "txt_attn.norm.query_norm.scale",
+                        "attn.norm_added_k.weight": "txt_attn.norm.key_norm.scale",
                     }
 
         for k in block_map:
@@ -449,15 +493,41 @@ def flux_to_diffusers(mmdit_config, output_prefix=""):
             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))
-            key_map["{}proj_mlp.{}".format(k, end)] = (qkv, (0, hidden_size * 3, hidden_size))
+            key_map["{}.proj_mlp.{}".format(prefix_from, end)] = (qkv, (0, hidden_size * 3, hidden_size * 4))
 
-        block_map = {#TODO
+        block_map = {
+                        "norm.linear.weight": "modulation.lin.weight",
+                        "norm.linear.bias": "modulation.lin.bias",
+                        "proj_out.weight": "linear2.weight",
+                        "proj_out.bias": "linear2.bias",
+                        "attn.norm_q.weight": "norm.query_norm.scale",
+                        "attn.norm_k.weight": "norm.key_norm.scale",
                     }
 
         for k in block_map:
             key_map["{}.{}".format(prefix_from, k)] = "{}.{}".format(prefix_to, block_map[k])
 
-    MAP_BASIC = { #TODO
+    MAP_BASIC = {
+        ("final_layer.linear.bias", "proj_out.bias"),
+        ("final_layer.linear.weight", "proj_out.weight"),
+        ("img_in.bias", "x_embedder.bias"),
+        ("img_in.weight", "x_embedder.weight"),
+        ("time_in.in_layer.bias", "time_text_embed.timestep_embedder.linear_1.bias"),
+        ("time_in.in_layer.weight", "time_text_embed.timestep_embedder.linear_1.weight"),
+        ("time_in.out_layer.bias", "time_text_embed.timestep_embedder.linear_2.bias"),
+        ("time_in.out_layer.weight", "time_text_embed.timestep_embedder.linear_2.weight"),
+        ("txt_in.bias", "context_embedder.bias"),
+        ("txt_in.weight", "context_embedder.weight"),
+        ("vector_in.in_layer.bias", "time_text_embed.text_embedder.linear_1.bias"),
+        ("vector_in.in_layer.weight", "time_text_embed.text_embedder.linear_1.weight"),
+        ("vector_in.out_layer.bias", "time_text_embed.text_embedder.linear_2.bias"),
+        ("vector_in.out_layer.weight", "time_text_embed.text_embedder.linear_2.weight"),
+        ("guidance_in.in_layer.bias", "time_text_embed.guidance_embedder.linear_1.bias"),
+        ("guidance_in.in_layer.weight", "time_text_embed.guidance_embedder.linear_1.weight"),
+        ("guidance_in.out_layer.bias", "time_text_embed.guidance_embedder.linear_2.bias"),
+        ("guidance_in.out_layer.weight", "time_text_embed.guidance_embedder.linear_2.weight"),
+        ("final_layer.adaLN_modulation.1.bias", "norm_out.linear.bias", swap_scale_shift),
+        ("final_layer.adaLN_modulation.1.weight", "norm_out.linear.weight", swap_scale_shift),
     }
 
     for k in MAP_BASIC:

comfy_execution/caching.py

@@ -0,0 +1,299 @@
+import itertools
+from typing import Sequence, Mapping
+from comfy_execution.graph import DynamicPrompt
+
+import nodes
+
+from comfy_execution.graph_utils import is_link
+
+class CacheKeySet:
+    def __init__(self, dynprompt, node_ids, is_changed_cache):
+        self.keys = {}
+        self.subcache_keys = {}
+
+    def add_keys(self, node_ids):
+        raise NotImplementedError()
+
+    def all_node_ids(self):
+        return set(self.keys.keys())
+
+    def get_used_keys(self):
+        return self.keys.values()
+
+    def get_used_subcache_keys(self):
+        return self.subcache_keys.values()
+
+    def get_data_key(self, node_id):
+        return self.keys.get(node_id, None)
+
+    def get_subcache_key(self, node_id):
+        return self.subcache_keys.get(node_id, None)
+
+class Unhashable:
+    def __init__(self):
+        self.value = float("NaN")
+
+def to_hashable(obj):
+    # So that we don't infinitely recurse since frozenset and tuples
+    # are Sequences.
+    if isinstance(obj, (int, float, str, bool, type(None))):
+        return obj
+    elif isinstance(obj, Mapping):
+        return frozenset([(to_hashable(k), to_hashable(v)) for k, v in sorted(obj.items())])
+    elif isinstance(obj, Sequence):
+        return frozenset(zip(itertools.count(), [to_hashable(i) for i in obj]))
+    else:
+        # TODO - Support other objects like tensors?
+        return Unhashable()
+
+class CacheKeySetID(CacheKeySet):
+    def __init__(self, dynprompt, node_ids, is_changed_cache):
+        super().__init__(dynprompt, node_ids, is_changed_cache)
+        self.dynprompt = dynprompt
+        self.add_keys(node_ids)
+
+    def add_keys(self, node_ids):
+        for node_id in node_ids:
+            if node_id in self.keys:
+                continue
+            node = self.dynprompt.get_node(node_id)
+            self.keys[node_id] = (node_id, node["class_type"])
+            self.subcache_keys[node_id] = (node_id, node["class_type"])
+
+class CacheKeySetInputSignature(CacheKeySet):
+    def __init__(self, dynprompt, node_ids, is_changed_cache):
+        super().__init__(dynprompt, node_ids, is_changed_cache)
+        self.dynprompt = dynprompt
+        self.is_changed_cache = is_changed_cache
+        self.add_keys(node_ids)
+
+    def include_node_id_in_input(self) -> bool:
+        return False
+
+    def add_keys(self, node_ids):
+        for node_id in node_ids:
+            if node_id in self.keys:
+                continue
+            node = self.dynprompt.get_node(node_id)
+            self.keys[node_id] = self.get_node_signature(self.dynprompt, node_id)
+            self.subcache_keys[node_id] = (node_id, node["class_type"])
+
+    def get_node_signature(self, dynprompt, node_id):
+        signature = []
+        ancestors, order_mapping = self.get_ordered_ancestry(dynprompt, node_id)
+        signature.append(self.get_immediate_node_signature(dynprompt, node_id, order_mapping))
+        for ancestor_id in ancestors:
+            signature.append(self.get_immediate_node_signature(dynprompt, ancestor_id, order_mapping))
+        return to_hashable(signature)
+
+    def get_immediate_node_signature(self, dynprompt, node_id, ancestor_order_mapping):
+        node = dynprompt.get_node(node_id)
+        class_type = node["class_type"]
+        class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+        signature = [class_type, self.is_changed_cache.get(node_id)]
+        if self.include_node_id_in_input() or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT):
+            signature.append(node_id)
+        inputs = node["inputs"]
+        for key in sorted(inputs.keys()):
+            if is_link(inputs[key]):
+                (ancestor_id, ancestor_socket) = inputs[key]
+                ancestor_index = ancestor_order_mapping[ancestor_id]
+                signature.append((key,("ANCESTOR", ancestor_index, ancestor_socket)))
+            else:
+                signature.append((key, inputs[key]))
+        return signature
+
+    # This function returns a list of all ancestors of the given node. The order of the list is
+    # deterministic based on which specific inputs the ancestor is connected by.
+    def get_ordered_ancestry(self, dynprompt, node_id):
+        ancestors = []
+        order_mapping = {}
+        self.get_ordered_ancestry_internal(dynprompt, node_id, ancestors, order_mapping)
+        return ancestors, order_mapping
+
+    def get_ordered_ancestry_internal(self, dynprompt, node_id, ancestors, order_mapping):
+        inputs = dynprompt.get_node(node_id)["inputs"]
+        input_keys = sorted(inputs.keys())
+        for key in input_keys:
+            if is_link(inputs[key]):
+                ancestor_id = inputs[key][0]
+                if ancestor_id not in order_mapping:
+                    ancestors.append(ancestor_id)
+                    order_mapping[ancestor_id] = len(ancestors) - 1
+                    self.get_ordered_ancestry_internal(dynprompt, ancestor_id, ancestors, order_mapping)
+
+class BasicCache:
+    def __init__(self, key_class):
+        self.key_class = key_class
+        self.initialized = False
+        self.dynprompt: DynamicPrompt
+        self.cache_key_set: CacheKeySet
+        self.cache = {}
+        self.subcaches = {}
+
+    def set_prompt(self, dynprompt, node_ids, is_changed_cache):
+        self.dynprompt = dynprompt
+        self.cache_key_set = self.key_class(dynprompt, node_ids, is_changed_cache)
+        self.is_changed_cache = is_changed_cache
+        self.initialized = True
+
+    def all_node_ids(self):
+        assert self.initialized
+        node_ids = self.cache_key_set.all_node_ids()
+        for subcache in self.subcaches.values():
+            node_ids = node_ids.union(subcache.all_node_ids())
+        return node_ids
+
+    def _clean_cache(self):
+        preserve_keys = set(self.cache_key_set.get_used_keys())
+        to_remove = []
+        for key in self.cache:
+            if key not in preserve_keys:
+                to_remove.append(key)
+        for key in to_remove:
+            del self.cache[key]
+
+    def _clean_subcaches(self):
+        preserve_subcaches = set(self.cache_key_set.get_used_subcache_keys())
+
+        to_remove = []
+        for key in self.subcaches:
+            if key not in preserve_subcaches:
+                to_remove.append(key)
+        for key in to_remove:
+            del self.subcaches[key]
+
+    def clean_unused(self):
+        assert self.initialized
+        self._clean_cache()
+        self._clean_subcaches()
+
+    def _set_immediate(self, node_id, value):
+        assert self.initialized
+        cache_key = self.cache_key_set.get_data_key(node_id)
+        self.cache[cache_key] = value
+
+    def _get_immediate(self, node_id):
+        if not self.initialized:
+            return None
+        cache_key = self.cache_key_set.get_data_key(node_id)
+        if cache_key in self.cache:
+            return self.cache[cache_key]
+        else:
+            return None
+
+    def _ensure_subcache(self, node_id, children_ids):
+        subcache_key = self.cache_key_set.get_subcache_key(node_id)
+        subcache = self.subcaches.get(subcache_key, None)
+        if subcache is None:
+            subcache = BasicCache(self.key_class)
+            self.subcaches[subcache_key] = subcache
+        subcache.set_prompt(self.dynprompt, children_ids, self.is_changed_cache)
+        return subcache
+
+    def _get_subcache(self, node_id):
+        assert self.initialized
+        subcache_key = self.cache_key_set.get_subcache_key(node_id)
+        if subcache_key in self.subcaches:
+            return self.subcaches[subcache_key]
+        else:
+            return None
+
+    def recursive_debug_dump(self):
+        result = []
+        for key in self.cache:
+            result.append({"key": key, "value": self.cache[key]})
+        for key in self.subcaches:
+            result.append({"subcache_key": key, "subcache": self.subcaches[key].recursive_debug_dump()})
+        return result
+
+class HierarchicalCache(BasicCache):
+    def __init__(self, key_class):
+        super().__init__(key_class)
+
+    def _get_cache_for(self, node_id):
+        assert self.dynprompt is not None
+        parent_id = self.dynprompt.get_parent_node_id(node_id)
+        if parent_id is None:
+            return self
+
+        hierarchy = []
+        while parent_id is not None:
+            hierarchy.append(parent_id)
+            parent_id = self.dynprompt.get_parent_node_id(parent_id)
+
+        cache = self
+        for parent_id in reversed(hierarchy):
+            cache = cache._get_subcache(parent_id)
+            if cache is None:
+                return None
+        return cache
+
+    def get(self, node_id):
+        cache = self._get_cache_for(node_id)
+        if cache is None:
+            return None
+        return cache._get_immediate(node_id)
+
+    def set(self, node_id, value):
+        cache = self._get_cache_for(node_id)
+        assert cache is not None
+        cache._set_immediate(node_id, value)
+
+    def ensure_subcache_for(self, node_id, children_ids):
+        cache = self._get_cache_for(node_id)
+        assert cache is not None
+        return cache._ensure_subcache(node_id, children_ids)
+
+class LRUCache(BasicCache):
+    def __init__(self, key_class, max_size=100):
+        super().__init__(key_class)
+        self.max_size = max_size
+        self.min_generation = 0
+        self.generation = 0
+        self.used_generation = {}
+        self.children = {}
+
+    def set_prompt(self, dynprompt, node_ids, is_changed_cache):
+        super().set_prompt(dynprompt, node_ids, is_changed_cache)
+        self.generation += 1
+        for node_id in node_ids:
+            self._mark_used(node_id)
+
+    def clean_unused(self):
+        while len(self.cache) > self.max_size and self.min_generation < self.generation:
+            self.min_generation += 1
+            to_remove = [key for key in self.cache if self.used_generation[key] < self.min_generation]
+            for key in to_remove:
+                del self.cache[key]
+                del self.used_generation[key]
+                if key in self.children:
+                    del self.children[key]
+        self._clean_subcaches()
+
+    def get(self, node_id):
+        self._mark_used(node_id)
+        return self._get_immediate(node_id)
+
+    def _mark_used(self, node_id):
+        cache_key = self.cache_key_set.get_data_key(node_id)
+        if cache_key is not None:
+            self.used_generation[cache_key] = self.generation
+
+    def set(self, node_id, value):
+        self._mark_used(node_id)
+        return self._set_immediate(node_id, value)
+
+    def ensure_subcache_for(self, node_id, children_ids):
+        # Just uses subcaches for tracking 'live' nodes
+        super()._ensure_subcache(node_id, children_ids)
+
+        self.cache_key_set.add_keys(children_ids)
+        self._mark_used(node_id)
+        cache_key = self.cache_key_set.get_data_key(node_id)
+        self.children[cache_key] = []
+        for child_id in children_ids:
+            self._mark_used(child_id)
+            self.children[cache_key].append(self.cache_key_set.get_data_key(child_id))
+        return self
+

comfy_execution/graph.py

@@ -0,0 +1,237 @@
+import nodes
+
+from comfy_execution.graph_utils import is_link
+
+class DependencyCycleError(Exception):
+    pass
+
+class NodeInputError(Exception):
+    pass
+
+class NodeNotFoundError(Exception):
+    pass
+
+class DynamicPrompt:
+    def __init__(self, original_prompt):
+        # The original prompt provided by the user
+        self.original_prompt = original_prompt
+        # Any extra pieces of the graph created during execution
+        self.ephemeral_prompt = {}
+        self.ephemeral_parents = {}
+        self.ephemeral_display = {}
+
+    def get_node(self, node_id):
+        if node_id in self.ephemeral_prompt:
+            return self.ephemeral_prompt[node_id]
+        if node_id in self.original_prompt:
+            return self.original_prompt[node_id]
+        raise NodeNotFoundError(f"Node {node_id} not found")
+
+    def has_node(self, node_id):
+        return node_id in self.original_prompt or node_id in self.ephemeral_prompt
+
+    def add_ephemeral_node(self, node_id, node_info, parent_id, display_id):
+        self.ephemeral_prompt[node_id] = node_info
+        self.ephemeral_parents[node_id] = parent_id
+        self.ephemeral_display[node_id] = display_id
+
+    def get_real_node_id(self, node_id):
+        while node_id in self.ephemeral_parents:
+            node_id = self.ephemeral_parents[node_id]
+        return node_id
+
+    def get_parent_node_id(self, node_id):
+        return self.ephemeral_parents.get(node_id, None)
+
+    def get_display_node_id(self, node_id):
+        while node_id in self.ephemeral_display:
+            node_id = self.ephemeral_display[node_id]
+        return node_id
+
+    def all_node_ids(self):
+        return set(self.original_prompt.keys()).union(set(self.ephemeral_prompt.keys()))
+
+    def get_original_prompt(self):
+        return self.original_prompt
+
+def get_input_info(class_def, input_name):
+    valid_inputs = class_def.INPUT_TYPES()
+    input_info = None
+    input_category = None
+    if "required" in valid_inputs and input_name in valid_inputs["required"]:
+        input_category = "required"
+        input_info = valid_inputs["required"][input_name]
+    elif "optional" in valid_inputs and input_name in valid_inputs["optional"]:
+        input_category = "optional"
+        input_info = valid_inputs["optional"][input_name]
+    elif "hidden" in valid_inputs and input_name in valid_inputs["hidden"]:
+        input_category = "hidden"
+        input_info = valid_inputs["hidden"][input_name]
+    if input_info is None:
+        return None, None, None
+    input_type = input_info[0]
+    if len(input_info) > 1:
+        extra_info = input_info[1]
+    else:
+        extra_info = {}
+    return input_type, input_category, extra_info
+
+class TopologicalSort:
+    def __init__(self, dynprompt):
+        self.dynprompt = dynprompt
+        self.pendingNodes = {}
+        self.blockCount = {} # Number of nodes this node is directly blocked by
+        self.blocking = {} # Which nodes are blocked by this node
+
+    def get_input_info(self, unique_id, input_name):
+        class_type = self.dynprompt.get_node(unique_id)["class_type"]
+        class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+        return get_input_info(class_def, input_name)
+
+    def make_input_strong_link(self, to_node_id, to_input):
+        inputs = self.dynprompt.get_node(to_node_id)["inputs"]
+        if to_input not in inputs:
+            raise NodeInputError(f"Node {to_node_id} says it needs input {to_input}, but there is no input to that node at all")
+        value = inputs[to_input]
+        if not is_link(value):
+            raise NodeInputError(f"Node {to_node_id} says it needs input {to_input}, but that value is a constant")
+        from_node_id, from_socket = value
+        self.add_strong_link(from_node_id, from_socket, to_node_id)
+
+    def add_strong_link(self, from_node_id, from_socket, to_node_id):
+        self.add_node(from_node_id)
+        if to_node_id not in self.blocking[from_node_id]:
+            self.blocking[from_node_id][to_node_id] = {}
+            self.blockCount[to_node_id] += 1
+        self.blocking[from_node_id][to_node_id][from_socket] = True
+
+    def add_node(self, unique_id, include_lazy=False, subgraph_nodes=None):
+        if unique_id in self.pendingNodes:
+            return
+        self.pendingNodes[unique_id] = True
+        self.blockCount[unique_id] = 0
+        self.blocking[unique_id] = {}
+
+        inputs = self.dynprompt.get_node(unique_id)["inputs"]
+        for input_name in inputs:
+            value = inputs[input_name]
+            if is_link(value):
+                from_node_id, from_socket = value
+                if subgraph_nodes is not None and from_node_id not in subgraph_nodes:
+                    continue
+                input_type, input_category, input_info = self.get_input_info(unique_id, input_name)
+                is_lazy = input_info is not None and "lazy" in input_info and input_info["lazy"]
+                if include_lazy or not is_lazy:
+                    self.add_strong_link(from_node_id, from_socket, unique_id)
+
+    def get_ready_nodes(self):
+        return [node_id for node_id in self.pendingNodes if self.blockCount[node_id] == 0]
+
+    def pop_node(self, unique_id):
+        del self.pendingNodes[unique_id]
+        for blocked_node_id in self.blocking[unique_id]:
+            self.blockCount[blocked_node_id] -= 1
+        del self.blocking[unique_id]
+
+    def is_empty(self):
+        return len(self.pendingNodes) == 0
+
+class ExecutionList(TopologicalSort):
+    """
+    ExecutionList implements a topological dissolve of the graph. After a node is staged for execution,
+    it can still be returned to the graph after having further dependencies added.
+    """
+    def __init__(self, dynprompt, output_cache):
+        super().__init__(dynprompt)
+        self.output_cache = output_cache
+        self.staged_node_id = None
+
+    def add_strong_link(self, from_node_id, from_socket, to_node_id):
+        if self.output_cache.get(from_node_id) is not None:
+            # Nothing to do
+            return
+        super().add_strong_link(from_node_id, from_socket, to_node_id)
+
+    def stage_node_execution(self):
+        assert self.staged_node_id is None
+        if self.is_empty():
+            return None, None, None
+        available = self.get_ready_nodes()
+        if len(available) == 0:
+            cycled_nodes = self.get_nodes_in_cycle()
+            # Because cycles composed entirely of static nodes are caught during initial validation,
+            # we will 'blame' the first node in the cycle that is not a static node.
+            blamed_node = cycled_nodes[0]
+            for node_id in cycled_nodes:
+                display_node_id = self.dynprompt.get_display_node_id(node_id)
+                if display_node_id != node_id:
+                    blamed_node = display_node_id
+                    break
+            ex = DependencyCycleError("Dependency cycle detected")
+            error_details = {
+                "node_id": blamed_node,
+                "exception_message": str(ex),
+                "exception_type": "graph.DependencyCycleError",
+                "traceback": [],
+                "current_inputs": []
+            }
+            return None, error_details, ex
+        next_node = available[0]
+        # If an output node is available, do that first.
+        # Technically this has no effect on the overall length of execution, but it feels better as a user
+        # for a PreviewImage to display a result as soon as it can
+        # Some other heuristics could probably be used here to improve the UX further.
+        for node_id in available:
+            class_type = self.dynprompt.get_node(node_id)["class_type"]
+            class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+            if hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE == True:
+                next_node = node_id
+                break
+        self.staged_node_id = next_node
+        return self.staged_node_id, None, None
+
+    def unstage_node_execution(self):
+        assert self.staged_node_id is not None
+        self.staged_node_id = None
+
+    def complete_node_execution(self):
+        node_id = self.staged_node_id
+        self.pop_node(node_id)
+        self.staged_node_id = None
+
+    def get_nodes_in_cycle(self):
+        # We'll dissolve the graph in reverse topological order to leave only the nodes in the cycle.
+        # We're skipping some of the performance optimizations from the original TopologicalSort to keep
+        # the code simple (and because having a cycle in the first place is a catastrophic error)
+        blocked_by = { node_id: {} for node_id in self.pendingNodes }
+        for from_node_id in self.blocking:
+            for to_node_id in self.blocking[from_node_id]:
+                if True in self.blocking[from_node_id][to_node_id].values():
+                    blocked_by[to_node_id][from_node_id] = True
+        to_remove = [node_id for node_id in blocked_by if len(blocked_by[node_id]) == 0]
+        while len(to_remove) > 0:
+            for node_id in to_remove:
+                for to_node_id in blocked_by:
+                    if node_id in blocked_by[to_node_id]:
+                        del blocked_by[to_node_id][node_id]
+                del blocked_by[node_id]
+            to_remove = [node_id for node_id in blocked_by if len(blocked_by[node_id]) == 0]
+        return list(blocked_by.keys())
+
+class ExecutionBlocker:
+    """
+    Return this from a node and any users will be blocked with the given error message.
+    If the message is None, execution will be blocked silently instead.
+    Generally, you should avoid using this functionality unless absolutely necessary. Whenever it's
+    possible, a lazy input will be more efficient and have a better user experience.
+    This functionality is useful in two cases:
+    1. You want to conditionally prevent an output node from executing. (Particularly a built-in node
+       like SaveImage. For your own output nodes, I would recommend just adding a BOOL input and using
+       lazy evaluation to let it conditionally disable itself.)
+    2. You have a node with multiple possible outputs, some of which are invalid and should not be used.
+       (I would recommend not making nodes like this in the future -- instead, make multiple nodes with
+       different outputs. Unfortunately, there are several popular existing nodes using this pattern.)
+    """
+    def __init__(self, message):
+        self.message = message
+

comfy_execution/graph_utils.py

@@ -0,0 +1,139 @@
+def is_link(obj):
+    if not isinstance(obj, list):
+        return False
+    if len(obj) != 2:
+        return False
+    if not isinstance(obj[0], str):
+        return False
+    if not isinstance(obj[1], int) and not isinstance(obj[1], float):
+        return False
+    return True
+
+# The GraphBuilder is just a utility class that outputs graphs in the form expected by the ComfyUI back-end
+class GraphBuilder:
+    _default_prefix_root = ""
+    _default_prefix_call_index = 0
+    _default_prefix_graph_index = 0
+
+    def __init__(self, prefix = None):
+        if prefix is None:
+            self.prefix = GraphBuilder.alloc_prefix()
+        else:
+            self.prefix = prefix
+        self.nodes = {}
+        self.id_gen = 1
+
+    @classmethod
+    def set_default_prefix(cls, prefix_root, call_index, graph_index = 0):
+        cls._default_prefix_root = prefix_root
+        cls._default_prefix_call_index = call_index
+        cls._default_prefix_graph_index = graph_index
+
+    @classmethod
+    def alloc_prefix(cls, root=None, call_index=None, graph_index=None):
+        if root is None:
+            root = GraphBuilder._default_prefix_root
+        if call_index is None:
+            call_index = GraphBuilder._default_prefix_call_index
+        if graph_index is None:
+            graph_index = GraphBuilder._default_prefix_graph_index
+        result = f"{root}.{call_index}.{graph_index}."
+        GraphBuilder._default_prefix_graph_index += 1
+        return result
+
+    def node(self, class_type, id=None, **kwargs):
+        if id is None:
+            id = str(self.id_gen)
+            self.id_gen += 1
+        id = self.prefix + id
+        if id in self.nodes:
+            return self.nodes[id]
+
+        node = Node(id, class_type, kwargs)
+        self.nodes[id] = node
+        return node
+
+    def lookup_node(self, id):
+        id = self.prefix + id
+        return self.nodes.get(id)
+
+    def finalize(self):
+        output = {}
+        for node_id, node in self.nodes.items():
+            output[node_id] = node.serialize()
+        return output
+
+    def replace_node_output(self, node_id, index, new_value):
+        node_id = self.prefix + node_id
+        to_remove = []
+        for node in self.nodes.values():
+            for key, value in node.inputs.items():
+                if is_link(value) and value[0] == node_id and value[1] == index:
+                    if new_value is None:
+                        to_remove.append((node, key))
+                    else:
+                        node.inputs[key] = new_value
+        for node, key in to_remove:
+            del node.inputs[key]
+
+    def remove_node(self, id):
+        id = self.prefix + id
+        del self.nodes[id]
+
+class Node:
+    def __init__(self, id, class_type, inputs):
+        self.id = id
+        self.class_type = class_type
+        self.inputs = inputs
+        self.override_display_id = None
+
+    def out(self, index):
+        return [self.id, index]
+
+    def set_input(self, key, value):
+        if value is None:
+            if key in self.inputs:
+                del self.inputs[key]
+        else:
+            self.inputs[key] = value
+
+    def get_input(self, key):
+        return self.inputs.get(key)
+
+    def set_override_display_id(self, override_display_id):
+        self.override_display_id = override_display_id
+
+    def serialize(self):
+        serialized = {
+            "class_type": self.class_type,
+            "inputs": self.inputs
+        }
+        if self.override_display_id is not None:
+            serialized["override_display_id"] = self.override_display_id
+        return serialized
+
+def add_graph_prefix(graph, outputs, prefix):
+    # Change the node IDs and any internal links
+    new_graph = {}
+    for node_id, node_info in graph.items():
+        # Make sure the added nodes have unique IDs
+        new_node_id = prefix + node_id
+        new_node = { "class_type": node_info["class_type"], "inputs": {} }
+        for input_name, input_value in node_info.get("inputs", {}).items():
+            if is_link(input_value):
+                new_node["inputs"][input_name] = [prefix + input_value[0], input_value[1]]
+            else:
+                new_node["inputs"][input_name] = input_value
+        new_graph[new_node_id] = new_node
+
+    # Change the node IDs in the outputs
+    new_outputs = []
+    for n in range(len(outputs)):
+        output = outputs[n]
+        if is_link(output):
+            new_outputs.append([prefix + output[0], output[1]])
+        else:
+            new_outputs.append(output)
+
+    return new_graph, tuple(new_outputs)
+

comfy_extras/nodes_hunyuan.py

@@ -19,6 +19,7 @@ class CLIPTextEncodeHunyuanDiT:
         cond = output.pop("cond")
         return ([[cond, output]], )
 
+
 NODE_CLASS_MAPPINGS = {
     "CLIPTextEncodeHunyuanDiT": CLIPTextEncodeHunyuanDiT,
 }

comfy_extras/nodes_sd3.py

@@ -100,3 +100,8 @@ NODE_CLASS_MAPPINGS = {
     "CLIPTextEncodeSD3": CLIPTextEncodeSD3,
     "ControlNetApplySD3": ControlNetApplySD3,
 }
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    # Sampling
+    "ControlNetApplySD3": "ControlNetApply SD3 and HunyuanDiT",
+}

custom_nodes/example_node.py.example

@@ -54,7 +54,8 @@ class Example:
                     "min": 0, #Minimum value
                     "max": 4096, #Maximum value
                     "step": 64, #Slider's step
-                    "display": "number" # Cosmetic only: display as "number" or "slider"
+                    "display": "number", # Cosmetic only: display as "number" or "slider"
+                    "lazy": True # Will only be evaluated if check_lazy_status requires it
                 }),
                 "float_field": ("FLOAT", {
                     "default": 1.0,
@@ -62,11 +63,14 @@ class Example:
                     "max": 10.0,
                     "step": 0.01,
                     "round": 0.001, #The value representing the precision to round to, will be set to the step value by default. Can be set to False to disable rounding.
-                    "display": "number"}),
+                    "display": "number",
+                    "lazy": True
+                }),
                 "print_to_screen": (["enable", "disable"],),
                 "string_field": ("STRING", {
                     "multiline": False, #True if you want the field to look like the one on the ClipTextEncode node
-                    "default": "Hello World!"
+                    "default": "Hello World!",
+                    "lazy": True
                 }),
             },
         }
@@ -80,6 +84,23 @@ class Example:
 
     CATEGORY = "Example"
 
+    def check_lazy_status(self, image, string_field, int_field, float_field, print_to_screen):
+        """
+            Return a list of input names that need to be evaluated.
+
+            This function will be called if there are any lazy inputs which have not yet been
+            evaluated. As long as you return at least one field which has not yet been evaluated
+            (and more exist), this function will be called again once the value of the requested
+            field is available.
+
+            Any evaluated inputs will be passed as arguments to this function. Any unevaluated
+            inputs will have the value None.
+        """
+        if print_to_screen == "enable":
+            return ["int_field", "float_field", "string_field"]
+        else:
+            return []
+
     def test(self, image, string_field, int_field, float_field, print_to_screen):
         if print_to_screen == "enable":
             print(f"""Your input contains:

execution.py

@@ -5,6 +5,7 @@ import threading
 import heapq
 import time
 import traceback
+from enum import Enum
 import inspect
 from typing import List, Literal, NamedTuple, Optional
 
@@ -12,102 +13,215 @@ import torch
 import nodes
 
 import comfy.model_management
+from comfy_execution.graph import get_input_info, ExecutionList, DynamicPrompt, ExecutionBlocker
+from comfy_execution.graph_utils import is_link, GraphBuilder
+from comfy_execution.caching import HierarchicalCache, LRUCache, CacheKeySetInputSignature, CacheKeySetID
+from comfy.cli_args import args
 
-def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}):
+class ExecutionResult(Enum):
+    SUCCESS = 0
+    FAILURE = 1
+    PENDING = 2
+
+class DuplicateNodeError(Exception):
+    pass
+
+class IsChangedCache:
+    def __init__(self, dynprompt, outputs_cache):
+        self.dynprompt = dynprompt
+        self.outputs_cache = outputs_cache
+        self.is_changed = {}
+
+    def get(self, node_id):
+        if node_id in self.is_changed:
+            return self.is_changed[node_id]
+
+        node = self.dynprompt.get_node(node_id)
+        class_type = node["class_type"]
+        class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+        if not hasattr(class_def, "IS_CHANGED"):
+            self.is_changed[node_id] = False
+            return self.is_changed[node_id]
+
+        if "is_changed" in node:
+            self.is_changed[node_id] = node["is_changed"]
+            return self.is_changed[node_id]
+
+        input_data_all, _ = get_input_data(node["inputs"], class_def, node_id, self.outputs_cache)
+        try:
+            is_changed = map_node_over_list(class_def, input_data_all, "IS_CHANGED")
+            node["is_changed"] = [None if isinstance(x, ExecutionBlocker) else x for x in is_changed]
+        except:
+            node["is_changed"] = float("NaN")
+        finally:
+            self.is_changed[node_id] = node["is_changed"]
+        return self.is_changed[node_id]
+
+class CacheSet:
+    def __init__(self, lru_size=None):
+        if lru_size is None or lru_size == 0:
+            self.init_classic_cache() 
+        else:
+            self.init_lru_cache(lru_size)
+        self.all = [self.outputs, self.ui, self.objects]
+
+    # Useful for those with ample RAM/VRAM -- allows experimenting without
+    # blowing away the cache every time
+    def init_lru_cache(self, cache_size):
+        self.outputs = LRUCache(CacheKeySetInputSignature, max_size=cache_size)
+        self.ui = LRUCache(CacheKeySetInputSignature, max_size=cache_size)
+        self.objects = HierarchicalCache(CacheKeySetID)
+
+    # Performs like the old cache -- dump data ASAP
+    def init_classic_cache(self):
+        self.outputs = HierarchicalCache(CacheKeySetInputSignature)
+        self.ui = HierarchicalCache(CacheKeySetInputSignature)
+        self.objects = HierarchicalCache(CacheKeySetID)
+
+    def recursive_debug_dump(self):
+        result = {
+            "outputs": self.outputs.recursive_debug_dump(),
+            "ui": self.ui.recursive_debug_dump(),
+        }
+        return result
+
+def get_input_data(inputs, class_def, unique_id, outputs=None, dynprompt=None, extra_data={}):
     valid_inputs = class_def.INPUT_TYPES()
     input_data_all = {}
+    missing_keys = {}
     for x in inputs:
         input_data = inputs[x]
-        if isinstance(input_data, list):
+        input_type, input_category, input_info = get_input_info(class_def, x)
+        def mark_missing():
+            missing_keys[x] = True
+            input_data_all[x] = (None,)
+        if is_link(input_data) and (not input_info or not input_info.get("rawLink", False)):
             input_unique_id = input_data[0]
             output_index = input_data[1]
-            if input_unique_id not in outputs:
-                input_data_all[x] = (None,)
+            if outputs is None:
+                mark_missing()
+                continue # This might be a lazily-evaluated input
+            cached_output = outputs.get(input_unique_id)
+            if cached_output is None:
+                mark_missing()
                 continue
-            obj = outputs[input_unique_id][output_index]
+            if output_index >= len(cached_output):
+                mark_missing()
+                continue
+            obj = cached_output[output_index]
             input_data_all[x] = obj
-        else:
-            if ("required" in valid_inputs and x in valid_inputs["required"]) or ("optional" in valid_inputs and x in valid_inputs["optional"]):
-                input_data_all[x] = [input_data]
+        elif input_category is not None:
+            input_data_all[x] = [input_data]
 
     if "hidden" in valid_inputs:
         h = valid_inputs["hidden"]
         for x in h:
             if h[x] == "PROMPT":
-                input_data_all[x] = [prompt]
+                input_data_all[x] = [dynprompt.get_original_prompt() if dynprompt is not None else {}]
+            if h[x] == "DYNPROMPT":
+                input_data_all[x] = [dynprompt]
             if h[x] == "EXTRA_PNGINFO":
                 input_data_all[x] = [extra_data.get('extra_pnginfo', None)]
             if h[x] == "UNIQUE_ID":
                 input_data_all[x] = [unique_id]
-    return input_data_all
+    return input_data_all, missing_keys
 
-def map_node_over_list(obj, input_data_all, func, allow_interrupt=False):
+def map_node_over_list(obj, input_data_all, func, allow_interrupt=False, execution_block_cb=None, pre_execute_cb=None):
     # check if node wants the lists
-    input_is_list = False
-    if hasattr(obj, "INPUT_IS_LIST"):
-        input_is_list = obj.INPUT_IS_LIST
+    input_is_list = getattr(obj, "INPUT_IS_LIST", False)
 
     if len(input_data_all) == 0:
         max_len_input = 0
     else:
-        max_len_input = max([len(x) for x in input_data_all.values()])
+        max_len_input = max(len(x) for x in input_data_all.values())
      
     # get a slice of inputs, repeat last input when list isn't long enough
     def slice_dict(d, i):
-        d_new = dict()
-        for k,v in d.items():
-            d_new[k] = v[i if len(v) > i else -1]
-        return d_new
+        return {k: v[i if len(v) > i else -1] for k, v in d.items()}
     
     results = []
+    def process_inputs(inputs, index=None):
+        if allow_interrupt:
+            nodes.before_node_execution()
+        execution_block = None
+        for k, v in inputs.items():
+            if isinstance(v, ExecutionBlocker):
+                execution_block = execution_block_cb(v) if execution_block_cb else v
+                break
+        if execution_block is None:
+            if pre_execute_cb is not None and index is not None:
+                pre_execute_cb(index)
+            results.append(getattr(obj, func)(**inputs))
+        else:
+            results.append(execution_block)
+
     if input_is_list:
-        if allow_interrupt:
-            nodes.before_node_execution()
-        results.append(getattr(obj, func)(**input_data_all))
+        process_inputs(input_data_all, 0)
     elif max_len_input == 0:
-        if allow_interrupt:
-            nodes.before_node_execution()
-        results.append(getattr(obj, func)())
-    else:
+        process_inputs({})
+    else: 
         for i in range(max_len_input):
-            if allow_interrupt:
-                nodes.before_node_execution()
-            results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
+            input_dict = slice_dict(input_data_all, i)
+            process_inputs(input_dict, i)
     return results
 
-def get_output_data(obj, input_data_all):
+def merge_result_data(results, obj):
+    # check which outputs need concatenating
+    output = []
+    output_is_list = [False] * len(results[0])
+    if hasattr(obj, "OUTPUT_IS_LIST"):
+        output_is_list = obj.OUTPUT_IS_LIST
+
+    # merge node execution results
+    for i, is_list in zip(range(len(results[0])), output_is_list):
+        if is_list:
+            output.append([x for o in results for x in o[i]])
+        else:
+            output.append([o[i] for o in results])
+    return output
+
+def get_output_data(obj, input_data_all, execution_block_cb=None, pre_execute_cb=None):
     
     results = []
     uis = []
-    return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)
-
-    for r in return_values:
+    subgraph_results = []
+    return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
+    has_subgraph = False
+    for i in range(len(return_values)):
+        r = return_values[i]
         if isinstance(r, dict):
             if 'ui' in r:
                 uis.append(r['ui'])
-            if 'result' in r:
-                results.append(r['result'])
+            if 'expand' in r:
+                # Perform an expansion, but do not append results
+                has_subgraph = True
+                new_graph = r['expand']
+                result = r.get("result", None)
+                if isinstance(result, ExecutionBlocker):
+                    result = tuple([result] * len(obj.RETURN_TYPES))
+                subgraph_results.append((new_graph, result))
+            elif 'result' in r:
+                result = r.get("result", None)
+                if isinstance(result, ExecutionBlocker):
+                    result = tuple([result] * len(obj.RETURN_TYPES))
+                results.append(result)
+                subgraph_results.append((None, result))
         else:
+            if isinstance(r, ExecutionBlocker):
+                r = tuple([r] * len(obj.RETURN_TYPES))
             results.append(r)
+            subgraph_results.append((None, r))
     
-    output = []
-    if len(results) > 0:
-        # check which outputs need concatenating
-        output_is_list = [False] * len(results[0])
-        if hasattr(obj, "OUTPUT_IS_LIST"):
-            output_is_list = obj.OUTPUT_IS_LIST
-
-        # merge node execution results
-        for i, is_list in zip(range(len(results[0])), output_is_list):
-            if is_list:
-                output.append([x for o in results for x in o[i]])
-            else:
-                output.append([o[i] for o in results])
-
+    if has_subgraph:
+        output = subgraph_results
+    elif len(results) > 0:
+        output = merge_result_data(results, obj)
+    else:
+        output = []
     ui = dict()    
     if len(uis) > 0:
         ui = {k: [y for x in uis for y in x[k]] for k in uis[0].keys()}
-    return output, ui
+    return output, ui, has_subgraph
 
 def format_value(x):
     if x is None:
@@ -117,53 +231,145 @@ def format_value(x):
     else:
         return str(x)
 
-def recursive_execute(server, prompt, outputs, current_item, extra_data, executed, prompt_id, outputs_ui, object_storage):
+def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results):
     unique_id = current_item
-    inputs = prompt[unique_id]['inputs']
-    class_type = prompt[unique_id]['class_type']
+    real_node_id = dynprompt.get_real_node_id(unique_id)
+    display_node_id = dynprompt.get_display_node_id(unique_id)
+    parent_node_id = dynprompt.get_parent_node_id(unique_id)
+    inputs = dynprompt.get_node(unique_id)['inputs']
+    class_type = dynprompt.get_node(unique_id)['class_type']
     class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
-    if unique_id in outputs:
-        return (True, None, None)
-
-    for x in inputs:
-        input_data = inputs[x]
-
-        if isinstance(input_data, list):
-            input_unique_id = input_data[0]
-            output_index = input_data[1]
-            if input_unique_id not in outputs:
-                result = recursive_execute(server, prompt, outputs, input_unique_id, extra_data, executed, prompt_id, outputs_ui, object_storage)
-                if result[0] is not True:
-                    # Another node failed further upstream
-                    return result
+    if caches.outputs.get(unique_id) is not None:
+        if server.client_id is not None:
+            cached_output = caches.ui.get(unique_id) or {}
+            server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": cached_output.get("output",None), "prompt_id": prompt_id }, server.client_id)
+        return (ExecutionResult.SUCCESS, None, None)
 
     input_data_all = None
     try:
-        input_data_all = get_input_data(inputs, class_def, unique_id, outputs, prompt, extra_data)
-        if server.client_id is not None:
-            server.last_node_id = unique_id
-            server.send_sync("executing", { "node": unique_id, "prompt_id": prompt_id }, server.client_id)
+        if unique_id in pending_subgraph_results:
+            cached_results = pending_subgraph_results[unique_id]
+            resolved_outputs = []
+            for is_subgraph, result in cached_results:
+                if not is_subgraph:
+                    resolved_outputs.append(result)
+                else:
+                    resolved_output = []
+                    for r in result:
+                        if is_link(r):
+                            source_node, source_output = r[0], r[1]
+                            node_output = caches.outputs.get(source_node)[source_output]
+                            for o in node_output:
+                                resolved_output.append(o)
 
-        obj = object_storage.get((unique_id, class_type), None)
-        if obj is None:
-            obj = class_def()
-            object_storage[(unique_id, class_type)] = obj
-
-        output_data, output_ui = get_output_data(obj, input_data_all)
-        outputs[unique_id] = output_data
-        if len(output_ui) > 0:
-            outputs_ui[unique_id] = output_ui
+                        else:
+                            resolved_output.append(r)
+                    resolved_outputs.append(tuple(resolved_output))
+            output_data = merge_result_data(resolved_outputs, class_def)
+            output_ui = []
+            has_subgraph = False
+        else:
+            input_data_all, missing_keys = get_input_data(inputs, class_def, unique_id, caches.outputs, dynprompt, extra_data)
             if server.client_id is not None:
-                server.send_sync("executed", { "node": unique_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id)
+                server.last_node_id = display_node_id
+                server.send_sync("executing", { "node": unique_id, "display_node": display_node_id, "prompt_id": prompt_id }, server.client_id)
+
+            obj = caches.objects.get(unique_id)
+            if obj is None:
+                obj = class_def()
+                caches.objects.set(unique_id, obj)
+
+            if hasattr(obj, "check_lazy_status"):
+                required_inputs = map_node_over_list(obj, input_data_all, "check_lazy_status", allow_interrupt=True)
+                required_inputs = set(sum([r for r in required_inputs if isinstance(r,list)], []))
+                required_inputs = [x for x in required_inputs if isinstance(x,str) and (
+                    x not in input_data_all or x in missing_keys
+                )]
+                if len(required_inputs) > 0:
+                    for i in required_inputs:
+                        execution_list.make_input_strong_link(unique_id, i)
+                    return (ExecutionResult.PENDING, None, None)
+
+            def execution_block_cb(block):
+                if block.message is not None:
+                    mes = {
+                        "prompt_id": prompt_id,
+                        "node_id": unique_id,
+                        "node_type": class_type,
+                        "executed": list(executed),
+
+                        "exception_message": f"Execution Blocked: {block.message}",
+                        "exception_type": "ExecutionBlocked",
+                        "traceback": [],
+                        "current_inputs": [],
+                        "current_outputs": [],
+                    }
+                    server.send_sync("execution_error", mes, server.client_id)
+                    return ExecutionBlocker(None)
+                else:
+                    return block
+            def pre_execute_cb(call_index):
+                GraphBuilder.set_default_prefix(unique_id, call_index, 0)
+            output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
+        if len(output_ui) > 0:
+            caches.ui.set(unique_id, {
+                "meta": {
+                    "node_id": unique_id,
+                    "display_node": display_node_id,
+                    "parent_node": parent_node_id,
+                    "real_node_id": real_node_id,
+                },
+                "output": output_ui
+            })
+            if server.client_id is not None:
+                server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id)
+        if has_subgraph:
+            cached_outputs = []
+            new_node_ids = []
+            new_output_ids = []
+            new_output_links = []
+            for i in range(len(output_data)):
+                new_graph, node_outputs = output_data[i]
+                if new_graph is None:
+                    cached_outputs.append((False, node_outputs))
+                else:
+                    # Check for conflicts
+                    for node_id in new_graph.keys():
+                        if dynprompt.has_node(node_id):
+                            raise DuplicateNodeError(f"Attempt to add duplicate node {node_id}. Ensure node ids are unique and deterministic or use graph_utils.GraphBuilder.")
+                    for node_id, node_info in new_graph.items():
+                        new_node_ids.append(node_id)
+                        display_id = node_info.get("override_display_id", unique_id)
+                        dynprompt.add_ephemeral_node(node_id, node_info, unique_id, display_id)
+                        # Figure out if the newly created node is an output node
+                        class_type = node_info["class_type"]
+                        class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+                        if hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE == True:
+                            new_output_ids.append(node_id)
+                    for i in range(len(node_outputs)):
+                        if is_link(node_outputs[i]):
+                            from_node_id, from_socket = node_outputs[i][0], node_outputs[i][1]
+                            new_output_links.append((from_node_id, from_socket))
+                    cached_outputs.append((True, node_outputs))
+            new_node_ids = set(new_node_ids)
+            for cache in caches.all:
+                cache.ensure_subcache_for(unique_id, new_node_ids).clean_unused()
+            for node_id in new_output_ids:
+                execution_list.add_node(node_id)
+            for link in new_output_links:
+                execution_list.add_strong_link(link[0], link[1], unique_id)
+            pending_subgraph_results[unique_id] = cached_outputs
+            return (ExecutionResult.PENDING, None, None)
+        caches.outputs.set(unique_id, output_data)
     except comfy.model_management.InterruptProcessingException as iex:
         logging.info("Processing interrupted")
 
         # skip formatting inputs/outputs
         error_details = {
-            "node_id": unique_id,
+            "node_id": real_node_id,
         }
 
-        return (False, error_details, iex)
+        return (ExecutionResult.FAILURE, error_details, iex)
     except Exception as ex:
         typ, _, tb = sys.exc_info()
         exception_type = full_type_name(typ)
@@ -173,121 +379,36 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute
             for name, inputs in input_data_all.items():
                 input_data_formatted[name] = [format_value(x) for x in inputs]
 
-        output_data_formatted = {}
-        for node_id, node_outputs in outputs.items():
-            output_data_formatted[node_id] = [[format_value(x) for x in l] for l in node_outputs]
-
-        logging.error(f"!!! Exception during processing!!! {ex}")
+        logging.error(f"!!! Exception during processing !!! {ex}")
         logging.error(traceback.format_exc())
 
         error_details = {
-            "node_id": unique_id,
+            "node_id": real_node_id,
             "exception_message": str(ex),
             "exception_type": exception_type,
             "traceback": traceback.format_tb(tb),
-            "current_inputs": input_data_formatted,
-            "current_outputs": output_data_formatted
+            "current_inputs": input_data_formatted
         }
-
         if isinstance(ex, comfy.model_management.OOM_EXCEPTION):
             logging.error("Got an OOM, unloading all loaded models.")
             comfy.model_management.unload_all_models()
 
-        return (False, error_details, ex)
+        return (ExecutionResult.FAILURE, error_details, ex)
 
     executed.add(unique_id)
 
-    return (True, None, None)
-
-def recursive_will_execute(prompt, outputs, current_item, memo={}):
-    unique_id = current_item
-
-    if unique_id in memo:
-        return memo[unique_id]
-
-    inputs = prompt[unique_id]['inputs']
-    will_execute = []
-    if unique_id in outputs:
-        return []
-
-    for x in inputs:
-        input_data = inputs[x]
-        if isinstance(input_data, list):
-            input_unique_id = input_data[0]
-            output_index = input_data[1]
-            if input_unique_id not in outputs:
-                will_execute += recursive_will_execute(prompt, outputs, input_unique_id, memo)
-
-    memo[unique_id] = will_execute + [unique_id]
-    return memo[unique_id]
-
-def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item):
-    unique_id = current_item
-    inputs = prompt[unique_id]['inputs']
-    class_type = prompt[unique_id]['class_type']
-    class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
-
-    is_changed_old = ''
-    is_changed = ''
-    to_delete = False
-    if hasattr(class_def, 'IS_CHANGED'):
-        if unique_id in old_prompt and 'is_changed' in old_prompt[unique_id]:
-            is_changed_old = old_prompt[unique_id]['is_changed']
-        if 'is_changed' not in prompt[unique_id]:
-            input_data_all = get_input_data(inputs, class_def, unique_id, outputs)
-            if input_data_all is not None:
-                try:
-                    #is_changed = class_def.IS_CHANGED(**input_data_all)
-                    is_changed = map_node_over_list(class_def, input_data_all, "IS_CHANGED")
-                    prompt[unique_id]['is_changed'] = is_changed
-                except:
-                    to_delete = True
-        else:
-            is_changed = prompt[unique_id]['is_changed']
-
-    if unique_id not in outputs:
-        return True
-
-    if not to_delete:
-        if is_changed != is_changed_old:
-            to_delete = True
-        elif unique_id not in old_prompt:
-            to_delete = True
-        elif class_type != old_prompt[unique_id]['class_type']:
-            to_delete = True
-        elif inputs == old_prompt[unique_id]['inputs']:
-            for x in inputs:
-                input_data = inputs[x]
-
-                if isinstance(input_data, list):
-                    input_unique_id = input_data[0]
-                    output_index = input_data[1]
-                    if input_unique_id in outputs:
-                        to_delete = recursive_output_delete_if_changed(prompt, old_prompt, outputs, input_unique_id)
-                    else:
-                        to_delete = True
-                    if to_delete:
-                        break
-        else:
-            to_delete = True
-
-    if to_delete:
-        d = outputs.pop(unique_id)
-        del d
-    return to_delete
+    return (ExecutionResult.SUCCESS, None, None)
 
 class PromptExecutor:
-    def __init__(self, server):
+    def __init__(self, server, lru_size=None):
+        self.lru_size = lru_size
         self.server = server
         self.reset()
 
     def reset(self):
-        self.outputs = {}
-        self.object_storage = {}
-        self.outputs_ui = {}
+        self.caches = CacheSet(self.lru_size)
         self.status_messages = []
         self.success = True
-        self.old_prompt = {}
 
     def add_message(self, event, data: dict, broadcast: bool):
         data = {
@@ -318,27 +439,14 @@ class PromptExecutor:
                 "node_id": node_id,
                 "node_type": class_type,
                 "executed": list(executed),
-
                 "exception_message": error["exception_message"],
                 "exception_type": error["exception_type"],
                 "traceback": error["traceback"],
                 "current_inputs": error["current_inputs"],
-                "current_outputs": error["current_outputs"],
+                "current_outputs": list(current_outputs),
             }
             self.add_message("execution_error", mes, broadcast=False)
         
-        # Next, remove the subsequent outputs since they will not be executed
-        to_delete = []
-        for o in self.outputs:
-            if (o not in current_outputs) and (o not in executed):
-                to_delete += [o]
-                if o in self.old_prompt:
-                    d = self.old_prompt.pop(o)
-                    del d
-        for o in to_delete:
-            d = self.outputs.pop(o)
-            del d
-
     def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
         nodes.interrupt_processing(False)
 
@@ -351,65 +459,58 @@ class PromptExecutor:
         self.add_message("execution_start", { "prompt_id": prompt_id}, broadcast=False)
 
         with torch.inference_mode():
-            #delete cached outputs if nodes don't exist for them
-            to_delete = []
-            for o in self.outputs:
-                if o not in prompt:
-                    to_delete += [o]
-            for o in to_delete:
-                d = self.outputs.pop(o)
-                del d
-            to_delete = []
-            for o in self.object_storage:
-                if o[0] not in prompt:
-                    to_delete += [o]
-                else:
-                    p = prompt[o[0]]
-                    if o[1] != p['class_type']:
-                        to_delete += [o]
-            for o in to_delete:
-                d = self.object_storage.pop(o)
-                del d
+            dynamic_prompt = DynamicPrompt(prompt)
+            is_changed_cache = IsChangedCache(dynamic_prompt, self.caches.outputs)
+            for cache in self.caches.all:
+                cache.set_prompt(dynamic_prompt, prompt.keys(), is_changed_cache)
+                cache.clean_unused()
 
-            for x in prompt:
-                recursive_output_delete_if_changed(prompt, self.old_prompt, self.outputs, x)
-
-            current_outputs = set(self.outputs.keys())
-            for x in list(self.outputs_ui.keys()):
-                if x not in current_outputs:
-                    d = self.outputs_ui.pop(x)
-                    del d
+            cached_nodes = []
+            for node_id in prompt:
+                if self.caches.outputs.get(node_id) is not None:
+                    cached_nodes.append(node_id)
 
             comfy.model_management.cleanup_models(keep_clone_weights_loaded=True)
             self.add_message("execution_cached",
-                          { "nodes": list(current_outputs) , "prompt_id": prompt_id},
+                          { "nodes": cached_nodes, "prompt_id": prompt_id},
                           broadcast=False)
+            pending_subgraph_results = {}
             executed = set()
-            output_node_id = None
-            to_execute = []
-
+            execution_list = ExecutionList(dynamic_prompt, self.caches.outputs)
+            current_outputs = self.caches.outputs.all_node_ids()
             for node_id in list(execute_outputs):
-                to_execute += [(0, node_id)]
+                execution_list.add_node(node_id)
 
-            while len(to_execute) > 0:
-                #always execute the output that depends on the least amount of unexecuted nodes first
-                memo = {}
-                to_execute = sorted(list(map(lambda a: (len(recursive_will_execute(prompt, self.outputs, a[-1], memo)), a[-1]), to_execute)))
-                output_node_id = to_execute.pop(0)[-1]
-
-                # This call shouldn't raise anything if there's an error deep in
-                # the actual SD code, instead it will report the node where the
-                # error was raised
-                self.success, error, ex = recursive_execute(self.server, prompt, self.outputs, output_node_id, extra_data, executed, prompt_id, self.outputs_ui, self.object_storage)
-                if self.success is not True:
-                    self.handle_execution_error(prompt_id, prompt, current_outputs, executed, error, ex)
+            while not execution_list.is_empty():
+                node_id, error, ex = execution_list.stage_node_execution()
+                if error is not None:
+                    self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
                     break
+
+                result, error, ex = execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results)
+                if result == ExecutionResult.FAILURE:
+                    self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
+                    break
+                elif result == ExecutionResult.PENDING:
+                    execution_list.unstage_node_execution()
+                else: # result == ExecutionResult.SUCCESS:
+                    execution_list.complete_node_execution()
             else:
                 # Only execute when the while-loop ends without break
                 self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False)
 
-            for x in executed:
-                self.old_prompt[x] = copy.deepcopy(prompt[x])
+            ui_outputs = {}
+            meta_outputs = {}
+            all_node_ids = self.caches.ui.all_node_ids()
+            for node_id in all_node_ids:
+                ui_info = self.caches.ui.get(node_id)
+                if ui_info is not None:
+                    ui_outputs[node_id] = ui_info["output"]
+                    meta_outputs[node_id] = ui_info["meta"]
+            self.history_result = {
+                "outputs": ui_outputs,
+                "meta": meta_outputs,
+            }
             self.server.last_node_id = None
             if comfy.model_management.DISABLE_SMART_MEMORY:
                 comfy.model_management.unload_all_models()
@@ -426,31 +527,37 @@ def validate_inputs(prompt, item, validated):
     obj_class = nodes.NODE_CLASS_MAPPINGS[class_type]
 
     class_inputs = obj_class.INPUT_TYPES()
-    required_inputs = class_inputs['required']
+    valid_inputs = set(class_inputs.get('required',{})).union(set(class_inputs.get('optional',{})))
 
     errors = []
     valid = True
 
     validate_function_inputs = []
+    validate_has_kwargs = False
     if hasattr(obj_class, "VALIDATE_INPUTS"):
-        validate_function_inputs = inspect.getfullargspec(obj_class.VALIDATE_INPUTS).args
+        argspec = inspect.getfullargspec(obj_class.VALIDATE_INPUTS)
+        validate_function_inputs = argspec.args
+        validate_has_kwargs = argspec.varkw is not None
+    received_types = {}
 
-    for x in required_inputs:
+    for x in valid_inputs:
+        type_input, input_category, extra_info = get_input_info(obj_class, x)
+        assert extra_info is not None
         if x not in inputs:
-            error = {
-                "type": "required_input_missing",
-                "message": "Required input is missing",
-                "details": f"{x}",
-                "extra_info": {
-                    "input_name": x
+            if input_category == "required":
+                error = {
+                    "type": "required_input_missing",
+                    "message": "Required input is missing",
+                    "details": f"{x}",
+                    "extra_info": {
+                        "input_name": x
+                    }
                 }
-            }
-            errors.append(error)
+                errors.append(error)
             continue
 
         val = inputs[x]
-        info = required_inputs[x]
-        type_input = info[0]
+        info = (type_input, extra_info)
         if isinstance(val, list):
             if len(val) != 2:
                 error = {
@@ -469,8 +576,9 @@ def validate_inputs(prompt, item, validated):
             o_id = val[0]
             o_class_type = prompt[o_id]['class_type']
             r = nodes.NODE_CLASS_MAPPINGS[o_class_type].RETURN_TYPES
-            if r[val[1]] != type_input:
-                received_type = r[val[1]]
+            received_type = r[val[1]]
+            received_types[x] = received_type
+            if 'input_types' not in validate_function_inputs and received_type != type_input:
                 details = f"{x}, {received_type} != {type_input}"
                 error = {
                     "type": "return_type_mismatch",
@@ -521,6 +629,9 @@ def validate_inputs(prompt, item, validated):
                 if type_input == "STRING":
                     val = str(val)
                     inputs[x] = val
+                if type_input == "BOOLEAN":
+                    val = bool(val)
+                    inputs[x] = val
             except Exception as ex:
                 error = {
                     "type": "invalid_input_type",
@@ -536,11 +647,11 @@ def validate_inputs(prompt, item, validated):
                 errors.append(error)
                 continue
 
-            if len(info) > 1:
-                if "min" in info[1] and val < info[1]["min"]:
+            if x not in validate_function_inputs and not validate_has_kwargs:
+                if "min" in extra_info and val < extra_info["min"]:
                     error = {
                         "type": "value_smaller_than_min",
-                        "message": "Value {} smaller than min of {}".format(val, info[1]["min"]),
+                        "message": "Value {} smaller than min of {}".format(val, extra_info["min"]),
                         "details": f"{x}",
                         "extra_info": {
                             "input_name": x,
@@ -550,10 +661,10 @@ def validate_inputs(prompt, item, validated):
                     }
                     errors.append(error)
                     continue
-                if "max" in info[1] and val > info[1]["max"]:
+                if "max" in extra_info and val > extra_info["max"]:
                     error = {
                         "type": "value_bigger_than_max",
-                        "message": "Value {} bigger than max of {}".format(val, info[1]["max"]),
+                        "message": "Value {} bigger than max of {}".format(val, extra_info["max"]),
                         "details": f"{x}",
                         "extra_info": {
                             "input_name": x,
@@ -564,7 +675,6 @@ def validate_inputs(prompt, item, validated):
                     errors.append(error)
                     continue
 
-            if x not in validate_function_inputs:
                 if isinstance(type_input, list):
                     if val not in type_input:
                         input_config = info
@@ -591,18 +701,20 @@ def validate_inputs(prompt, item, validated):
                         errors.append(error)
                         continue
 
-    if len(validate_function_inputs) > 0:
-        input_data_all = get_input_data(inputs, obj_class, unique_id)
+    if len(validate_function_inputs) > 0 or validate_has_kwargs:
+        input_data_all, _ = get_input_data(inputs, obj_class, unique_id)
         input_filtered = {}
         for x in input_data_all:
-            if x in validate_function_inputs:
+            if x in validate_function_inputs or validate_has_kwargs:
                 input_filtered[x] = input_data_all[x]
+        if 'input_types' in validate_function_inputs:
+            input_filtered['input_types'] = [received_types]
 
         #ret = obj_class.VALIDATE_INPUTS(**input_filtered)
         ret = map_node_over_list(obj_class, input_filtered, "VALIDATE_INPUTS")
         for x in input_filtered:
             for i, r in enumerate(ret):
-                if r is not True:
+                if r is not True and not isinstance(r, ExecutionBlocker):
                     details = f"{x}"
                     if r is not False:
                         details += f" - {str(r)}"
@@ -613,8 +725,6 @@ def validate_inputs(prompt, item, validated):
                         "details": details,
                         "extra_info": {
                             "input_name": x,
-                            "input_config": info,
-                            "received_value": val,
                         }
                     }
                     errors.append(error)
@@ -780,7 +890,7 @@ class PromptQueue:
         completed: bool
         messages: List[str]
 
-    def task_done(self, item_id, outputs,
+    def task_done(self, item_id, history_result,
                   status: Optional['PromptQueue.ExecutionStatus']):
         with self.mutex:
             prompt = self.currently_running.pop(item_id)
@@ -793,9 +903,10 @@ class PromptQueue:
 
             self.history[prompt[1]] = {
                 "prompt": prompt,
-                "outputs": copy.deepcopy(outputs),
+                "outputs": {},
                 'status': status_dict,
             }
+            self.history[prompt[1]].update(history_result)
             self.server.queue_updated()
 
     def get_current_queue(self):

main.py

@@ -101,7 +101,7 @@ def cuda_malloc_warning():
             logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n")
 
 def prompt_worker(q, server):
-    e = execution.PromptExecutor(server)
+    e = execution.PromptExecutor(server, lru_size=args.cache_lru)
     last_gc_collect = 0
     need_gc = False
     gc_collect_interval = 10.0
@@ -121,7 +121,7 @@ def prompt_worker(q, server):
             e.execute(item[2], prompt_id, item[3], item[4])
             need_gc = True
             q.task_done(item_id,
-                        e.outputs_ui,
+                        e.history_result,
                         status=execution.PromptQueue.ExecutionStatus(
                             status_str='success' if e.success else 'error',
                             completed=e.success,
@@ -261,6 +261,7 @@ if __name__ == "__main__":
         call_on_start = startup_server
 
     try:
+        loop.run_until_complete(server.setup())
         loop.run_until_complete(run(server, address=args.listen, port=args.port, verbose=not args.dont_print_server, call_on_start=call_on_start))
     except KeyboardInterrupt:
         logging.info("\nStopped server")

model_filemanager/__init__.py

@@ -0,0 +1,2 @@
+# model_manager/__init__.py
+from .download_models import download_model, DownloadModelStatus, DownloadStatusType, create_model_path, check_file_exists, track_download_progress, validate_model_subdirectory, validate_filename

model_filemanager/download_models.py

@@ -0,0 +1,240 @@
+from __future__ import annotations
+import aiohttp
+import os
+import traceback
+import logging
+from folder_paths import models_dir
+import re
+from typing import Callable, Any, Optional, Awaitable, Dict
+from enum import Enum
+import time
+from dataclasses import dataclass
+
+
+class DownloadStatusType(Enum):
+    PENDING = "pending"
+    IN_PROGRESS = "in_progress"
+    COMPLETED = "completed"
+    ERROR = "error"
+
+@dataclass
+class DownloadModelStatus():
+    status: str
+    progress_percentage: float
+    message: str
+    already_existed: bool = False
+
+    def __init__(self, status: DownloadStatusType, progress_percentage: float, message: str, already_existed: bool):
+        self.status = status.value  # Store the string value of the Enum
+        self.progress_percentage = progress_percentage
+        self.message = message
+        self.already_existed = already_existed
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "status": self.status,
+            "progress_percentage": self.progress_percentage,
+            "message": self.message,
+            "already_existed": self.already_existed
+        }
+
+async def download_model(model_download_request: Callable[[str], Awaitable[aiohttp.ClientResponse]],
+                         model_name: str,  
+                         model_url: str, 
+                         model_sub_directory: str,
+                         progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]],
+                         progress_interval: float = 1.0) -> DownloadModelStatus:
+    """
+    Download a model file from a given URL into the models directory.
+
+    Args:
+        model_download_request (Callable[[str], Awaitable[aiohttp.ClientResponse]]): 
+            A function that makes an HTTP request. This makes it easier to mock in unit tests.
+        model_name (str): 
+            The name of the model file to be downloaded. This will be the filename on disk.
+        model_url (str): 
+            The URL from which to download the model.
+        model_sub_directory (str): 
+            The subdirectory within the main models directory where the model 
+            should be saved (e.g., 'checkpoints', 'loras', etc.).
+        progress_callback (Callable[[str, DownloadModelStatus], Awaitable[Any]]): 
+            An asynchronous function to call with progress updates.
+
+    Returns:
+        DownloadModelStatus: The result of the download operation.
+    """
+    if not validate_model_subdirectory(model_sub_directory):
+        return DownloadModelStatus(
+            DownloadStatusType.ERROR, 
+            0,
+            "Invalid model subdirectory", 
+            False
+        )
+
+    if not validate_filename(model_name):
+        return DownloadModelStatus(
+            DownloadStatusType.ERROR, 
+            0,
+            "Invalid model name", 
+            False
+        )
+
+    file_path, relative_path = create_model_path(model_name, model_sub_directory, models_dir)
+    existing_file = await check_file_exists(file_path, model_name, progress_callback, relative_path)
+    if existing_file:
+        return existing_file
+
+    try:
+        status = DownloadModelStatus(DownloadStatusType.PENDING, 0, f"Starting download of {model_name}", False)
+        await progress_callback(relative_path, status)
+
+        response = await model_download_request(model_url)
+        if response.status != 200:
+            error_message = f"Failed to download {model_name}. Status code: {response.status}"
+            logging.error(error_message)
+            status = DownloadModelStatus(DownloadStatusType.ERROR, 0, error_message, False)
+            await progress_callback(relative_path, status)
+            return DownloadModelStatus(DownloadStatusType.ERROR, 0, error_message, False)
+
+        return await track_download_progress(response, file_path, model_name, progress_callback, relative_path, progress_interval)
+
+    except Exception as e:
+        logging.error(f"Error in downloading model: {e}")
+        return await handle_download_error(e, model_name, progress_callback, relative_path)
+    
+
+def create_model_path(model_name: str, model_directory: str, models_base_dir: str) -> tuple[str, str]:
+    full_model_dir = os.path.join(models_base_dir, model_directory)
+    os.makedirs(full_model_dir, exist_ok=True)
+    file_path = os.path.join(full_model_dir, model_name)
+
+    # Ensure the resulting path is still within the base directory
+    abs_file_path = os.path.abspath(file_path)
+    abs_base_dir = os.path.abspath(str(models_base_dir))
+    if os.path.commonprefix([abs_file_path, abs_base_dir]) != abs_base_dir:
+        raise Exception(f"Invalid model directory: {model_directory}/{model_name}")
+
+
+    relative_path = '/'.join([model_directory, model_name])
+    return file_path, relative_path
+
+async def check_file_exists(file_path: str, 
+                            model_name: str, 
+                            progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]], 
+                            relative_path: str) -> Optional[DownloadModelStatus]:
+    if os.path.exists(file_path):
+        status = DownloadModelStatus(DownloadStatusType.COMPLETED, 100, f"{model_name} already exists", True)
+        await progress_callback(relative_path, status)
+        return status
+    return None
+
+
+async def track_download_progress(response: aiohttp.ClientResponse, 
+                                  file_path: str, 
+                                  model_name: str, 
+                                  progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]], 
+                                  relative_path: str, 
+                                  interval: float = 1.0) -> DownloadModelStatus:
+    try:
+        total_size = int(response.headers.get('Content-Length', 0))
+        downloaded = 0
+        last_update_time = time.time()
+
+        async def update_progress():
+            nonlocal last_update_time
+            progress = (downloaded / total_size) * 100 if total_size > 0 else 0
+            status = DownloadModelStatus(DownloadStatusType.IN_PROGRESS, progress, f"Downloading {model_name}", False)
+            await progress_callback(relative_path, status)
+            last_update_time = time.time()
+
+        with open(file_path, 'wb') as f:
+            chunk_iterator = response.content.iter_chunked(8192)
+            while True:
+                try:
+                    chunk = await chunk_iterator.__anext__()
+                except StopAsyncIteration:
+                    break
+                f.write(chunk)
+                downloaded += len(chunk)
+                
+                if time.time() - last_update_time >= interval:
+                    await update_progress()
+
+        await update_progress()
+        
+        logging.info(f"Successfully downloaded {model_name}. Total downloaded: {downloaded}")
+        status = DownloadModelStatus(DownloadStatusType.COMPLETED, 100, f"Successfully downloaded {model_name}", False)
+        await progress_callback(relative_path, status)
+
+        return status
+    except Exception as e:
+        logging.error(f"Error in track_download_progress: {e}")
+        logging.error(traceback.format_exc())
+        return await handle_download_error(e, model_name, progress_callback, relative_path)
+
+async def handle_download_error(e: Exception, 
+                                model_name: str, 
+                                progress_callback: Callable[[str, DownloadModelStatus], Any], 
+                                relative_path: str) -> DownloadModelStatus:
+    error_message = f"Error downloading {model_name}: {str(e)}"
+    status = DownloadModelStatus(DownloadStatusType.ERROR, 0, error_message, False)
+    await progress_callback(relative_path, status)
+    return status
+
+def validate_model_subdirectory(model_subdirectory: str) -> bool:
+    """
+    Validate that the model subdirectory is safe to install into. 
+    Must not contain relative paths, nested paths or special characters
+    other than underscores and hyphens.
+
+    Args:
+        model_subdirectory (str): The subdirectory for the specific model type.
+
+    Returns:
+        bool: True if the subdirectory is safe, False otherwise.
+    """
+    if len(model_subdirectory) > 50:
+        return False
+
+    if '..' in model_subdirectory or '/' in model_subdirectory:
+        return False
+
+    if not re.match(r'^[a-zA-Z0-9_-]+$', model_subdirectory):
+        return False
+
+    return True
+
+def validate_filename(filename: str)-> bool:
+    """
+    Validate a filename to ensure it's safe and doesn't contain any path traversal attempts.
+    
+    Args:
+    filename (str): The filename to validate
+
+    Returns:
+    bool: True if the filename is valid, False otherwise
+    """
+    if not filename.lower().endswith(('.sft', '.safetensors')):
+        return False
+
+    # Check if the filename is empty, None, or just whitespace
+    if not filename or not filename.strip():
+        return False
+
+    # Check for any directory traversal attempts or invalid characters
+    if any(char in filename for char in ['..', '/', '\\', '\n', '\r', '\t', '\0']):
+        return False
+
+    # Check if the filename starts with a dot (hidden file)
+    if filename.startswith('.'):
+        return False
+
+    # Use a whitelist of allowed characters
+    if not re.match(r'^[a-zA-Z0-9_\-. ]+$', filename):
+        return False
+
+    # Ensure the filename isn't too long
+    if len(filename) > 255:
+        return False
+
+    return True

nodes.py

@@ -47,11 +47,18 @@ MAX_RESOLUTION=16384
 class CLIPTextEncode:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": {"text": ("STRING", {"multiline": True, "dynamicPrompts": True}), "clip": ("CLIP", )}}
+        return {
+            "required": {
+                "text": ("STRING", {"multiline": True, "dynamicPrompts": True, "tooltip": "The text to be encoded."}), 
+                "clip": ("CLIP", {"tooltip": "The CLIP model used for encoding the text."})
+            }
+        }
     RETURN_TYPES = ("CONDITIONING",)
+    OUTPUT_TOOLTIPS = ("A conditioning containing the embedded text used to guide the diffusion model.",)
     FUNCTION = "encode"
 
     CATEGORY = "conditioning"
+    DESCRIPTION = "Encodes a text prompt using a CLIP model into an embedding that can be used to guide the diffusion model towards generating specific images."
 
     def encode(self, clip, text):
         tokens = clip.tokenize(text)
@@ -260,11 +267,18 @@ class ConditioningSetTimestepRange:
 class VAEDecode:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT", ), "vae": ("VAE", )}}
+        return {
+            "required": { 
+                "samples": ("LATENT", {"tooltip": "The latent to be decoded."}), 
+                "vae": ("VAE", {"tooltip": "The VAE model used for decoding the latent."})
+            }
+        }
     RETURN_TYPES = ("IMAGE",)
+    OUTPUT_TOOLTIPS = ("The decoded image.",)
     FUNCTION = "decode"
 
     CATEGORY = "latent"
+    DESCRIPTION = "Decodes latent images back into pixel space images."
 
     def decode(self, vae, samples):
         return (vae.decode(samples["samples"]), )
@@ -506,12 +520,19 @@ class CheckpointLoader:
 class CheckpointLoaderSimple:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
-                             }}
+        return {
+            "required": { 
+                "ckpt_name": (folder_paths.get_filename_list("checkpoints"), {"tooltip": "The name of the checkpoint (model) to load."}),
+            }
+        }
     RETURN_TYPES = ("MODEL", "CLIP", "VAE")
+    OUTPUT_TOOLTIPS = ("The model used for denoising latents.", 
+                       "The CLIP model used for encoding text prompts.", 
+                       "The VAE model used for encoding and decoding images to and from latent space.")
     FUNCTION = "load_checkpoint"
 
     CATEGORY = "loaders"
+    DESCRIPTION = "Loads a diffusion model checkpoint, diffusion models are used to denoise latents."
 
     def load_checkpoint(self, ckpt_name):
         ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
@@ -582,16 +603,22 @@ class LoraLoader:
 
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "clip": ("CLIP", ),
-                              "lora_name": (folder_paths.get_filename_list("loras"), ),
-                              "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
-                              "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
-                              }}
+        return {
+            "required": { 
+                "model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}),
+                "clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}),
+                "lora_name": (folder_paths.get_filename_list("loras"), {"tooltip": "The name of the LoRA."}),
+                "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the diffusion model. This value can be negative."}),
+                "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}),
+            }
+        }
+    
     RETURN_TYPES = ("MODEL", "CLIP")
+    OUTPUT_TOOLTIPS = ("The modified diffusion model.", "The modified CLIP model.")
     FUNCTION = "load_lora"
 
     CATEGORY = "loaders"
+    DESCRIPTION = "LoRAs are used to modify diffusion and CLIP models, altering the way in which latents are denoised such as applying styles. Multiple LoRA nodes can be linked together."
 
     def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
         if strength_model == 0 and strength_clip == 0:
@@ -826,14 +853,14 @@ class UNETLoader:
     CATEGORY = "advanced/loaders"
 
     def load_unet(self, unet_name, weight_dtype):
-        dtype = None
+        model_options = {}
         if weight_dtype == "fp8_e4m3fn":
-            dtype = torch.float8_e4m3fn
+            model_options["dtype"] = torch.float8_e4m3fn
         elif weight_dtype == "fp8_e5m2":
-            dtype = torch.float8_e5m2
+            model_options["dtype"] = torch.float8_e5m2
 
         unet_path = folder_paths.get_full_path("unet", unet_name)
-        model = comfy.sd.load_unet(unet_path, dtype=dtype)
+        model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
         return (model,)
 
 class CLIPLoader:
@@ -1033,13 +1060,19 @@ class EmptyLatentImage:
 
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": { "width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}}
+        return {
+            "required": { 
+                "width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The width of the latent images in pixels."}),
+                "height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The height of the latent images in pixels."}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."})
+            }
+        }
     RETURN_TYPES = ("LATENT",)
+    OUTPUT_TOOLTIPS = ("The empty latent image batch.",)
     FUNCTION = "generate"
 
     CATEGORY = "latent"
+    DESCRIPTION = "Create a new batch of empty latent images to be denoised via sampling."
 
     def generate(self, width, height, batch_size=1):
         latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
@@ -1359,24 +1392,27 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
 class KSampler:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
-                    "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
-                    "positive": ("CONDITIONING", ),
-                    "negative": ("CONDITIONING", ),
-                    "latent_image": ("LATENT", ),
-                    "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                     }
-                }
+        return {
+            "required": {
+                "model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}),
+                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}),
+                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."}),
+                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
+                "positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}),
+                "negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
+                "latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}),
+                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising applied, lower values will maintain the structure of the initial image allowing for image to image sampling."}),
+            }
+        }
 
     RETURN_TYPES = ("LATENT",)
+    OUTPUT_TOOLTIPS = ("The denoised latent.",)
     FUNCTION = "sample"
 
     CATEGORY = "sampling"
+    DESCRIPTION = "Uses the provided model, positive and negative conditioning to denoise the latent image."
 
     def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
         return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
@@ -1424,11 +1460,15 @@ class SaveImage:
 
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": 
-                    {"images": ("IMAGE", ),
-                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
-                }
+        return {
+            "required": {
+                "images": ("IMAGE", {"tooltip": "The images to save."}),
+                "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
+            },
+            "hidden": {
+                "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
+            },
+        }
 
     RETURN_TYPES = ()
     FUNCTION = "save_images"
@@ -1436,6 +1476,7 @@ class SaveImage:
     OUTPUT_NODE = True
 
     CATEGORY = "image"
+    DESCRIPTION = "Saves the input images to your ComfyUI output directory."
 
     def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, output_dir=None):
         # ===========================

pytest.ini

@@ -1,6 +1,7 @@
 [pytest]
 markers = 
   inference: mark as inference test (deselect with '-m "not inference"')
+  execution: mark as execution test (deselect with '-m "not execution"')
 testpaths =
   tests
   tests-unit

server.py

@@ -12,7 +12,6 @@ import json
 import glob
 import struct
 import ssl
-import hashlib
 from PIL import Image, ImageOps
 from PIL.PngImagePlugin import PngInfo
 from io import BytesIO
@@ -28,7 +27,8 @@ import comfy.model_management
 import node_helpers
 from app.frontend_management import FrontendManager
 from app.user_manager import UserManager
-
+from model_filemanager import download_model, DownloadModelStatus
+from typing import Optional
 
 class BinaryEventTypes:
     PREVIEW_IMAGE = 1
@@ -76,6 +76,7 @@ class PromptServer():
         self.prompt_queue = None
         self.loop = loop
         self.messages = asyncio.Queue()
+        self.client_session:Optional[aiohttp.ClientSession] = None
         self.number = 0
 
         middlewares = [cache_control]
@@ -422,6 +423,7 @@ class PromptServer():
             obj_class = nodes.NODE_CLASS_MAPPINGS[node_class]
             info = {}
             info['input'] = obj_class.INPUT_TYPES()
+            info['input_order'] = {key: list(value.keys()) for (key, value) in obj_class.INPUT_TYPES().items()}
             info['output'] = obj_class.RETURN_TYPES
             info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES)
             info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output']
@@ -437,6 +439,9 @@ class PromptServer():
 
             if hasattr(obj_class, 'CATEGORY'):
                 info['category'] = obj_class.CATEGORY
+
+            if hasattr(obj_class, 'OUTPUT_TOOLTIPS'):
+                info['output_tooltips'] = obj_class.OUTPUT_TOOLTIPS
             return info
 
         @routes.get("/object_info")
@@ -559,6 +564,36 @@ class PromptServer():
                     self.prompt_queue.delete_history_item(id_to_delete)
 
             return web.Response(status=200)
+        
+        # Internal route. Should not be depended upon and is subject to change at any time.
+        # TODO(robinhuang): Move to internal route table class once we refactor PromptServer to pass around Websocket.
+        @routes.post("/internal/models/download")
+        async def download_handler(request):
+            async def report_progress(filename: str, status: DownloadModelStatus):
+                await self.send_json("download_progress", status.to_dict())
+
+            data = await request.json()
+            url = data.get('url')
+            model_directory = data.get('model_directory')
+            model_filename = data.get('model_filename')
+            progress_interval = data.get('progress_interval', 1.0) # In seconds, how often to report download progress.
+
+            if not url or not model_directory or not model_filename:
+                return web.json_response({"status": "error", "message": "Missing URL or folder path or filename"}, status=400)
+
+            session = self.client_session
+            if session is None:
+                logging.error("Client session is not initialized")
+                return web.Response(status=500)
+            
+            task = asyncio.create_task(download_model(lambda url: session.get(url), model_filename, url, model_directory, report_progress, progress_interval))
+            await task
+
+            return web.json_response(task.result().to_dict())
+
+    async def setup(self):
+        timeout = aiohttp.ClientTimeout(total=None) # no timeout
+        self.client_session = aiohttp.ClientSession(timeout=timeout)
 
     def add_routes(self):
         self.user_manager.add_routes(self.routes)
@@ -680,6 +715,9 @@ class PromptServer():
         site = web.TCPSite(runner, address, port, ssl_context=ssl_ctx)
         await site.start()
 
+        self.address = address
+        self.port = port
+
         if verbose:
             logging.info("Starting server\n")
             logging.info("To see the GUI go to: {}://{}:{}".format(scheme, address, port))

tests-ui/.gitignore

@@ -1 +0,0 @@
-node_modules

tests-ui/afterSetup.js

@@ -1,9 +0,0 @@
-const { start } = require("./utils");
-const lg = require("./utils/litegraph");
-
-// Load things once per test file before to ensure its all warmed up for the tests
-beforeAll(async () => {
-	lg.setup(global);
-	await start({ resetEnv: true });
-	lg.teardown(global);
-});

tests-ui/babel.config.json

@@ -1,4 +0,0 @@
-{
-	"presets": ["@babel/preset-env"],
-	"plugins": ["babel-plugin-transform-import-meta"]
-}

tests-ui/globalSetup.js

@@ -1,14 +0,0 @@
-module.exports = async function () {
-	global.ResizeObserver = class ResizeObserver {
-		observe() {}
-		unobserve() {}
-		disconnect() {}
-	};
-
-	const { nop } = require("./utils/nopProxy");
-	global.enableWebGLCanvas = nop;
-
-	HTMLCanvasElement.prototype.getContext = nop;
-
-	localStorage["Comfy.Settings.Comfy.Logging.Enabled"] = "false";
-};

tests-ui/jest.config.js

@@ -1,11 +0,0 @@
-/** @type {import('jest').Config} */
-const config = {
-	testEnvironment: "jsdom",
-	setupFiles: ["./globalSetup.js"],
-	setupFilesAfterEnv: ["./afterSetup.js"],
-	clearMocks: true,
-	resetModules: true,
-	testTimeout: 10000
-};
-
-module.exports = config;

tests-ui/package.json

@@ -1,31 +0,0 @@
-{
-  "name": "comfui-tests",
-  "version": "1.0.0",
-  "description": "UI tests",
-  "main": "index.js",
-  "scripts": {
-    "test": "jest",
-    "test:generate": "node setup.js"
-  },
-  "repository": {
-    "type": "git",
-    "url": "git+https://github.com/comfyanonymous/ComfyUI.git"
-  },
-  "keywords": [
-    "comfyui",
-    "test"
-  ],
-  "author": "comfyanonymous",
-  "license": "GPL-3.0",
-  "bugs": {
-    "url": "https://github.com/comfyanonymous/ComfyUI/issues"
-  },
-  "homepage": "https://github.com/comfyanonymous/ComfyUI#readme",
-  "devDependencies": {
-    "@babel/preset-env": "^7.22.20",
-    "@types/jest": "^29.5.5",
-    "babel-plugin-transform-import-meta": "^2.2.1",
-    "jest": "^29.7.0",
-    "jest-environment-jsdom": "^29.7.0"
-  }
-}

tests-ui/setup.js

@@ -1,88 +0,0 @@
-const { spawn } = require("child_process");
-const { resolve } = require("path");
-const { existsSync, mkdirSync, writeFileSync } = require("fs");
-const http = require("http");
-
-async function setup() {
-	// Wait up to 30s for it to start
-	let success = false;
-	let child;
-	for (let i = 0; i < 30; i++) {
-		try {
-			await new Promise((res, rej) => {
-				http
-					.get("http://127.0.0.1:8188/object_info", (resp) => {
-						let data = "";
-						resp.on("data", (chunk) => {
-							data += chunk;
-						});
-						resp.on("end", () => {
-							// Modify the response data to add some checkpoints
-							const objectInfo = JSON.parse(data);
-							objectInfo.CheckpointLoaderSimple.input.required.ckpt_name[0] = ["model1.safetensors", "model2.ckpt"];
-							objectInfo.VAELoader.input.required.vae_name[0] = ["vae1.safetensors", "vae2.ckpt"];
-
-							data = JSON.stringify(objectInfo, undefined, "\t");
-
-							const outDir = resolve("./data");
-							if (!existsSync(outDir)) {
-								mkdirSync(outDir);
-							}
-
-							const outPath = resolve(outDir, "object_info.json");
-							console.log(`Writing ${Object.keys(objectInfo).length} nodes to ${outPath}`);
-							writeFileSync(outPath, data, {
-								encoding: "utf8",
-							});
-							res();
-						});
-					})
-					.on("error", rej);
-			});
-			success = true;
-			break;
-		} catch (error) {
-			console.log(i + "/30", error);
-			if (i === 0) {
-				// Start the server on first iteration if it fails to connect
-				console.log("Starting ComfyUI server...");
-
-				let python = resolve("../../python_embeded/python.exe");
-				let args;
-				let cwd;
-				if (existsSync(python)) {
-					args = ["-s", "ComfyUI/main.py"];
-					cwd = "../..";
-				} else {
-					python = "python";
-					args = ["main.py"];
-					cwd = "..";
-				}
-				args.push("--cpu");
-				console.log(python, ...args);
-				child = spawn(python, args, { cwd });
-				child.on("error", (err) => {
-					console.log(`Server error (${err})`);
-					i = 30;
-				});
-				child.on("exit", (code) => {
-					if (!success) {
-						console.log(`Server exited (${code})`);
-						i = 30;
-					}
-				});
-			}
-			await new Promise((r) => {
-				setTimeout(r, 1000);
-			});
-		}
-	}
-
-	child?.kill();
-
-	if (!success) {
-		throw new Error("Waiting for server failed...");
-	}
-}
-
- setup();

tests-ui/tests/extensions.test.js

@@ -1,196 +0,0 @@
-// @ts-check
-/// <reference path="../node_modules/@types/jest/index.d.ts" />
-const { start } = require("../utils");
-const lg = require("../utils/litegraph");
-
-describe("extensions", () => {
-	beforeEach(() => {
-		lg.setup(global);
-	});
-
-	afterEach(() => {
-		lg.teardown(global);
-	});
-
-	it("calls each extension hook", async () => {
-		const mockExtension = {
-			name: "TestExtension",
-			init: jest.fn(),
-			setup: jest.fn(),
-			addCustomNodeDefs: jest.fn(),
-			getCustomWidgets: jest.fn(),
-			beforeRegisterNodeDef: jest.fn(),
-			registerCustomNodes: jest.fn(),
-			loadedGraphNode: jest.fn(),
-			nodeCreated: jest.fn(),
-			beforeConfigureGraph: jest.fn(),
-			afterConfigureGraph: jest.fn(),
-		};
-
-		const { app, ez, graph } = await start({
-			async preSetup(app) {
-				app.registerExtension(mockExtension);
-			},
-		});
-
-		// Basic initialisation hooks should be called once, with app
-		expect(mockExtension.init).toHaveBeenCalledTimes(1);
-		expect(mockExtension.init).toHaveBeenCalledWith(app);
-
-		// Adding custom node defs should be passed the full list of nodes
-		expect(mockExtension.addCustomNodeDefs).toHaveBeenCalledTimes(1);
-		expect(mockExtension.addCustomNodeDefs.mock.calls[0][1]).toStrictEqual(app);
-		const defs = mockExtension.addCustomNodeDefs.mock.calls[0][0];
-		expect(defs).toHaveProperty("KSampler");
-		expect(defs).toHaveProperty("LoadImage");
-
-		// Get custom widgets is called once and should return new widget types
-		expect(mockExtension.getCustomWidgets).toHaveBeenCalledTimes(1);
-		expect(mockExtension.getCustomWidgets).toHaveBeenCalledWith(app);
-
-		// Before register node def will be called once per node type
-		const nodeNames = Object.keys(defs);
-		const nodeCount = nodeNames.length;
-		expect(mockExtension.beforeRegisterNodeDef).toHaveBeenCalledTimes(nodeCount);
-		for (let i = 0; i < 10; i++) {
-			// It should be send the JS class and the original JSON definition
-			const nodeClass = mockExtension.beforeRegisterNodeDef.mock.calls[i][0];
-			const nodeDef = mockExtension.beforeRegisterNodeDef.mock.calls[i][1];
-
-			expect(nodeClass.name).toBe("ComfyNode");
-			expect(nodeClass.comfyClass).toBe(nodeNames[i]);
-			expect(nodeDef.name).toBe(nodeNames[i]);
-			expect(nodeDef).toHaveProperty("input");
-			expect(nodeDef).toHaveProperty("output");
-		}
-
-		// Register custom nodes is called once after registerNode defs to allow adding other frontend nodes
-		expect(mockExtension.registerCustomNodes).toHaveBeenCalledTimes(1);
-
-		// Before configure graph will be called here as the default graph is being loaded
-		expect(mockExtension.beforeConfigureGraph).toHaveBeenCalledTimes(1);
-		// it gets sent the graph data that is going to be loaded
-		const graphData = mockExtension.beforeConfigureGraph.mock.calls[0][0];
-
-		// A node created is fired for each node constructor that is called
-		expect(mockExtension.nodeCreated).toHaveBeenCalledTimes(graphData.nodes.length);
-		for (let i = 0; i < graphData.nodes.length; i++) {
-			expect(mockExtension.nodeCreated.mock.calls[i][0].type).toBe(graphData.nodes[i].type);
-		}
-
-		// Each node then calls loadedGraphNode to allow them to be updated
-		expect(mockExtension.loadedGraphNode).toHaveBeenCalledTimes(graphData.nodes.length);
-		for (let i = 0; i < graphData.nodes.length; i++) {
-			expect(mockExtension.loadedGraphNode.mock.calls[i][0].type).toBe(graphData.nodes[i].type);
-		}
-
-		// After configure is then called once all the setup is done
-		expect(mockExtension.afterConfigureGraph).toHaveBeenCalledTimes(1);
-
-		expect(mockExtension.setup).toHaveBeenCalledTimes(1);
-		expect(mockExtension.setup).toHaveBeenCalledWith(app);
-
-		// Ensure hooks are called in the correct order
-		const callOrder = [
-			"init",
-			"addCustomNodeDefs",
-			"getCustomWidgets",
-			"beforeRegisterNodeDef",
-			"registerCustomNodes",
-			"beforeConfigureGraph",
-			"nodeCreated",
-			"loadedGraphNode",
-			"afterConfigureGraph",
-			"setup",
-		];
-		for (let i = 1; i < callOrder.length; i++) {
-			const fn1 = mockExtension[callOrder[i - 1]];
-			const fn2 = mockExtension[callOrder[i]];
-			expect(fn1.mock.invocationCallOrder[0]).toBeLessThan(fn2.mock.invocationCallOrder[0]);
-		}
-
-		graph.clear();
-
-		// Ensure adding a new node calls the correct callback
-		ez.LoadImage();
-		expect(mockExtension.loadedGraphNode).toHaveBeenCalledTimes(graphData.nodes.length);
-		expect(mockExtension.nodeCreated).toHaveBeenCalledTimes(graphData.nodes.length + 1);
-		expect(mockExtension.nodeCreated.mock.lastCall[0].type).toBe("LoadImage");
-
-		// Reload the graph to ensure correct hooks are fired
-		await graph.reload();
-
-		// These hooks should not be fired again
-		expect(mockExtension.init).toHaveBeenCalledTimes(1);
-		expect(mockExtension.addCustomNodeDefs).toHaveBeenCalledTimes(1);
-		expect(mockExtension.getCustomWidgets).toHaveBeenCalledTimes(1);
-		expect(mockExtension.registerCustomNodes).toHaveBeenCalledTimes(1);
-		expect(mockExtension.beforeRegisterNodeDef).toHaveBeenCalledTimes(nodeCount);
-		expect(mockExtension.setup).toHaveBeenCalledTimes(1);
-
-		// These should be called again
-		expect(mockExtension.beforeConfigureGraph).toHaveBeenCalledTimes(2);
-		expect(mockExtension.nodeCreated).toHaveBeenCalledTimes(graphData.nodes.length + 2);
-		expect(mockExtension.loadedGraphNode).toHaveBeenCalledTimes(graphData.nodes.length + 1);
-		expect(mockExtension.afterConfigureGraph).toHaveBeenCalledTimes(2);
-	}, 15000);
-
-	it("allows custom nodeDefs and widgets to be registered", async () => {
-		const widgetMock = jest.fn((node, inputName, inputData, app) => {
-			expect(node.constructor.comfyClass).toBe("TestNode");
-			expect(inputName).toBe("test_input");
-			expect(inputData[0]).toBe("CUSTOMWIDGET");
-			expect(inputData[1]?.hello).toBe("world");
-			expect(app).toStrictEqual(app);
-
-			return {
-				widget: node.addWidget("button", inputName, "hello", () => {}),
-			};
-		});
-
-		// Register our extension that adds a custom node + widget type
-		const mockExtension = {
-			name: "TestExtension",
-			addCustomNodeDefs: (nodeDefs) => {
-				nodeDefs["TestNode"] = {
-					output: [],
-					output_name: [],
-					output_is_list: [],
-					name: "TestNode",
-					display_name: "TestNode",
-					category: "Test",
-					input: {
-						required: {
-							test_input: ["CUSTOMWIDGET", { hello: "world" }],
-						},
-					},
-				};
-			},
-			getCustomWidgets: jest.fn(() => {
-				return {
-					CUSTOMWIDGET: widgetMock,
-				};
-			}),
-		};
-
-		const { graph, ez } = await start({
-			async preSetup(app) {
-				app.registerExtension(mockExtension);
-			},
-		});
-
-		expect(mockExtension.getCustomWidgets).toBeCalledTimes(1);
-
-		graph.clear();
-		expect(widgetMock).toBeCalledTimes(0);
-		const node = ez.TestNode();
-		expect(widgetMock).toBeCalledTimes(1);
-
-		// Ensure our custom widget is created
-		expect(node.inputs.length).toBe(0);
-		expect(node.widgets.length).toBe(1);
-		const w = node.widgets[0].widget;
-		expect(w.name).toBe("test_input");
-		expect(w.type).toBe("button");
-	});
-});

tests-ui/tests/users.test.js

@@ -1,295 +0,0 @@
-// @ts-check
-/// <reference path="../node_modules/@types/jest/index.d.ts" />
-const { start } = require("../utils");
-const lg = require("../utils/litegraph");
-
-describe("users", () => {
-	beforeEach(() => {
-		lg.setup(global);
-	});
-
-	afterEach(() => {
-		lg.teardown(global);
-	});
-
-	function expectNoUserScreen() {
-		// Ensure login isnt visible
-		const selection = document.querySelectorAll("#comfy-user-selection")?.[0];
-		expect(selection["style"].display).toBe("none");
-		const menu = document.querySelectorAll(".comfy-menu")?.[0];
-		expect(window.getComputedStyle(menu)?.display).not.toBe("none");
-	}
-
-	describe("multi-user", () => {
-		function mockAddStylesheet() {
-			const utils = require("../../web/scripts/utils");
-			utils.addStylesheet = jest.fn().mockReturnValue(Promise.resolve());
-		}
-
-		async function waitForUserScreenShow() {
-			mockAddStylesheet();
-
-			// Wait for "show" to be called
-			const { UserSelectionScreen } = require("../../web/scripts/ui/userSelection");
-			let resolve, reject;
-			const fn = UserSelectionScreen.prototype.show;
-			const p = new Promise((res, rej) => {
-				resolve = res;
-				reject = rej;
-			});
-			jest.spyOn(UserSelectionScreen.prototype, "show").mockImplementation(async (...args) => {
-				const res = fn(...args);
-				await new Promise(process.nextTick); // wait for promises to resolve
-				resolve();
-				return res;
-			});
-			// @ts-ignore
-			setTimeout(() => reject("timeout waiting for UserSelectionScreen to be shown."), 500);
-			await p;
-			await new Promise(process.nextTick); // wait for promises to resolve
-		}
-
-		async function testUserScreen(onShown, users) {
-			if (!users) {
-				users = {};
-			}
-			const starting = start({
-				resetEnv: true,
-				userConfig: { storage: "server", users },
-			});
-
-			// Ensure no current user
-			expect(localStorage["Comfy.userId"]).toBeFalsy();
-			expect(localStorage["Comfy.userName"]).toBeFalsy();
-
-			await waitForUserScreenShow();
-
-			const selection = document.querySelectorAll("#comfy-user-selection")?.[0];
-			expect(selection).toBeTruthy();
-
-			// Ensure login is visible
-			expect(window.getComputedStyle(selection)?.display).not.toBe("none");
-			// Ensure menu is hidden
-			const menu = document.querySelectorAll(".comfy-menu")?.[0];
-			expect(window.getComputedStyle(menu)?.display).toBe("none");
-
-			const isCreate = await onShown(selection);
-
-			// Submit form
-			selection.querySelectorAll("form")[0].submit();
-			await new Promise(process.nextTick); // wait for promises to resolve
-
-			// Wait for start
-			const s = await starting;
-
-			// Ensure login is removed
-			expect(document.querySelectorAll("#comfy-user-selection")).toHaveLength(0);
-			expect(window.getComputedStyle(menu)?.display).not.toBe("none");
-
-			// Ensure settings + templates are saved
-			const { api } = require("../../web/scripts/api");
-			expect(api.createUser).toHaveBeenCalledTimes(+isCreate);
-			expect(api.storeSettings).toHaveBeenCalledTimes(+isCreate);
-			expect(api.storeUserData).toHaveBeenCalledTimes(+isCreate);
-			if (isCreate) {
-				expect(api.storeUserData).toHaveBeenCalledWith("comfy.templates.json", null, { stringify: false });
-				expect(s.app.isNewUserSession).toBeTruthy();
-			} else {
-				expect(s.app.isNewUserSession).toBeFalsy();
-			}
-
-			return { users, selection, ...s };
-		}
-
-		it("allows user creation if no users", async () => {
-			const { users } = await testUserScreen((selection) => {
-				// Ensure we have no users flag added
-				expect(selection.classList.contains("no-users")).toBeTruthy();
-
-				// Enter a username
-				const input = selection.getElementsByTagName("input")[0];
-				input.focus();
-				input.value = "Test User";
-
-				return true;
-			});
-
-			expect(users).toStrictEqual({
-				"Test User!": "Test User",
-			});
-
-			expect(localStorage["Comfy.userId"]).toBe("Test User!");
-			expect(localStorage["Comfy.userName"]).toBe("Test User");
-		});
-		it("allows user creation if no current user but other users", async () => {
-			const users = {
-				"Test User 2!": "Test User 2",
-			};
-
-			await testUserScreen((selection) => {
-				expect(selection.classList.contains("no-users")).toBeFalsy();
-
-				// Enter a username
-				const input = selection.getElementsByTagName("input")[0];
-				input.focus();
-				input.value = "Test User 3";
-				return true;
-			}, users);
-
-			expect(users).toStrictEqual({
-				"Test User 2!": "Test User 2",
-				"Test User 3!": "Test User 3",
-			});
-
-			expect(localStorage["Comfy.userId"]).toBe("Test User 3!");
-			expect(localStorage["Comfy.userName"]).toBe("Test User 3");
-		});
-		it("allows user selection if no current user but other users", async () => {
-			const users = {
-				"A!": "A",
-				"B!": "B",
-				"C!": "C",
-			};
-
-			await testUserScreen((selection) => {
-				expect(selection.classList.contains("no-users")).toBeFalsy();
-
-				// Check user list
-				const select = selection.getElementsByTagName("select")[0];
-				const options = select.getElementsByTagName("option");
-				expect(
-					[...options]
-						.filter((o) => !o.disabled)
-						.reduce((p, n) => {
-							p[n.getAttribute("value")] = n.textContent;
-							return p;
-						}, {})
-				).toStrictEqual(users);
-
-				// Select an option
-				select.focus();
-				select.value = options[2].value;
-
-				return false;
-			}, users);
-
-			expect(users).toStrictEqual(users);
-
-			expect(localStorage["Comfy.userId"]).toBe("B!");
-			expect(localStorage["Comfy.userName"]).toBe("B");
-		});
-		it("doesnt show user screen if current user", async () => {
-			const starting = start({
-				resetEnv: true,
-				userConfig: {
-					storage: "server",
-					users: {
-						"User!": "User",
-					},
-				},
-				localStorage: {
-					"Comfy.userId": "User!",
-					"Comfy.userName": "User",
-				},
-			});
-			await new Promise(process.nextTick); // wait for promises to resolve
-
-			expectNoUserScreen();
-
-			await starting;
-		});
-		it("allows user switching", async () => {
-			const { app } = await start({
-				resetEnv: true,
-				userConfig: {
-					storage: "server",
-					users: {
-						"User!": "User",
-					},
-				},
-				localStorage: {
-					"Comfy.userId": "User!",
-					"Comfy.userName": "User",
-				},
-			});
-
-			// cant actually test switching user easily but can check the setting is present
-			expect(app.ui.settings.settingsLookup["Comfy.SwitchUser"]).toBeTruthy();
-		});
-	});
-	describe("single-user", () => {
-		it("doesnt show user creation if no default user", async () => {
-			const { app } = await start({
-				resetEnv: true,
-				userConfig: { migrated: false, storage: "server" },
-			});
-			expectNoUserScreen();
-
-			// It should store the settings
-			const { api } = require("../../web/scripts/api");
-			expect(api.storeSettings).toHaveBeenCalledTimes(1);
-			expect(api.storeUserData).toHaveBeenCalledTimes(1);
-			expect(api.storeUserData).toHaveBeenCalledWith("comfy.templates.json", null, { stringify: false });
-			expect(app.isNewUserSession).toBeTruthy();
-		});
-		it("doesnt show user creation if default user", async () => {
-			const { app } = await start({
-				resetEnv: true,
-				userConfig: { migrated: true, storage: "server" },
-			});
-			expectNoUserScreen();
-
-			// It should store the settings
-			const { api } = require("../../web/scripts/api");
-			expect(api.storeSettings).toHaveBeenCalledTimes(0);
-			expect(api.storeUserData).toHaveBeenCalledTimes(0);
-			expect(app.isNewUserSession).toBeFalsy();
-		});
-		it("doesnt allow user switching", async () => {
-			const { app } = await start({
-				resetEnv: true,
-				userConfig: { migrated: true, storage: "server" },
-			});
-			expectNoUserScreen();
-
-			expect(app.ui.settings.settingsLookup["Comfy.SwitchUser"]).toBeFalsy();
-		});
-	});
-	describe("browser-user", () => {
-		it("doesnt show user creation if no default user", async () => {
-			const { app } = await start({
-				resetEnv: true,
-				userConfig: { migrated: false, storage: "browser" },
-			});
-			expectNoUserScreen();
-
-			// It should store the settings
-			const { api } = require("../../web/scripts/api");
-			expect(api.storeSettings).toHaveBeenCalledTimes(0);
-			expect(api.storeUserData).toHaveBeenCalledTimes(0);
-			expect(app.isNewUserSession).toBeFalsy();
-		});
-		it("doesnt show user creation if default user", async () => {
-			const { app } = await start({
-				resetEnv: true,
-				userConfig: { migrated: true, storage: "server" },
-			});
-			expectNoUserScreen();
-
-			// It should store the settings
-			const { api } = require("../../web/scripts/api");
-			expect(api.storeSettings).toHaveBeenCalledTimes(0);
-			expect(api.storeUserData).toHaveBeenCalledTimes(0);
-			expect(app.isNewUserSession).toBeFalsy();
-		});
-		it("doesnt allow user switching", async () => {
-			const { app } = await start({
-				resetEnv: true,
-				userConfig: { migrated: true, storage: "browser" },
-			});
-			expectNoUserScreen();
-
-			expect(app.ui.settings.settingsLookup["Comfy.SwitchUser"]).toBeFalsy();
-		});
-	});
-});

tests-ui/tests/widgetInputs.test.js

@@ -1,557 +0,0 @@
-// @ts-check
-/// <reference path="../node_modules/@types/jest/index.d.ts" />
-
-const {
-	start,
-	makeNodeDef,
-	checkBeforeAndAfterReload,
-	assertNotNullOrUndefined,
-	createDefaultWorkflow,
-} = require("../utils");
-const lg = require("../utils/litegraph");
-
-/**
- * @typedef { import("../utils/ezgraph") } Ez
- * @typedef { ReturnType<Ez["Ez"]["graph"]>["ez"] } EzNodeFactory
- */
-
-/**
- * @param { EzNodeFactory } ez
- * @param { InstanceType<Ez["EzGraph"]> } graph
- * @param { InstanceType<Ez["EzInput"]> } input
- * @param { string } widgetType
- * @param { number } controlWidgetCount
- * @returns
- */
-async function connectPrimitiveAndReload(ez, graph, input, widgetType, controlWidgetCount = 0) {
-	// Connect to primitive and ensure its still connected after
-	let primitive = ez.PrimitiveNode();
-	primitive.outputs[0].connectTo(input);
-
-	await checkBeforeAndAfterReload(graph, async () => {
-		primitive = graph.find(primitive);
-		let { connections } = primitive.outputs[0];
-		expect(connections).toHaveLength(1);
-		expect(connections[0].targetNode.id).toBe(input.node.node.id);
-
-		// Ensure widget is correct type
-		const valueWidget = primitive.widgets.value;
-		expect(valueWidget.widget.type).toBe(widgetType);
-
-		// Check if control_after_generate should be added
-		if (controlWidgetCount) {
-			const controlWidget = primitive.widgets.control_after_generate;
-			expect(controlWidget.widget.type).toBe("combo");
-			if (widgetType === "combo") {
-				const filterWidget = primitive.widgets.control_filter_list;
-				expect(filterWidget.widget.type).toBe("string");
-			}
-		}
-
-		// Ensure we dont have other widgets
-		expect(primitive.node.widgets).toHaveLength(1 + controlWidgetCount);
-	});
-
-	return primitive;
-}
-
-describe("widget inputs", () => {
-	beforeEach(() => {
-		lg.setup(global);
-	});
-
-	afterEach(() => {
-		lg.teardown(global);
-	});
-
-	[
-		{ name: "int", type: "INT", widget: "number", control: 1 },
-		{ name: "float", type: "FLOAT", widget: "number", control: 1 },
-		{ name: "text", type: "STRING" },
-		{
-			name: "customtext",
-			type: "STRING",
-			opt: { multiline: true },
-		},
-		{ name: "toggle", type: "BOOLEAN" },
-		{ name: "combo", type: ["a", "b", "c"], control: 2 },
-	].forEach((c) => {
-		test(`widget conversion + primitive works on ${c.name}`, async () => {
-			const { ez, graph } = await start({
-				mockNodeDefs: makeNodeDef("TestNode", { [c.name]: [c.type, c.opt ?? {}] }),
-			});
-
-			// Create test node and convert to input
-			const n = ez.TestNode();
-			const w = n.widgets[c.name];
-			w.convertToInput();
-			expect(w.isConvertedToInput).toBeTruthy();
-			const input = w.getConvertedInput();
-			expect(input).toBeTruthy();
-
-			// @ts-ignore : input is valid here
-			await connectPrimitiveAndReload(ez, graph, input, c.widget ?? c.name, c.control);
-		});
-	});
-
-	test("converted widget works after reload", async () => {
-		const { ez, graph } = await start();
-		let n = ez.CheckpointLoaderSimple();
-
-		const inputCount = n.inputs.length;
-
-		// Convert ckpt name to an input
-		n.widgets.ckpt_name.convertToInput();
-		expect(n.widgets.ckpt_name.isConvertedToInput).toBeTruthy();
-		expect(n.inputs.ckpt_name).toBeTruthy();
-		expect(n.inputs.length).toEqual(inputCount + 1);
-
-		// Convert back to widget and ensure input is removed
-		n.widgets.ckpt_name.convertToWidget();
-		expect(n.widgets.ckpt_name.isConvertedToInput).toBeFalsy();
-		expect(n.inputs.ckpt_name).toBeFalsy();
-		expect(n.inputs.length).toEqual(inputCount);
-
-		// Convert again and reload the graph to ensure it maintains state
-		n.widgets.ckpt_name.convertToInput();
-		expect(n.inputs.length).toEqual(inputCount + 1);
-
-		const primitive = await connectPrimitiveAndReload(ez, graph, n.inputs.ckpt_name, "combo", 2);
-
-		// Disconnect & reconnect
-		primitive.outputs[0].connections[0].disconnect();
-		let { connections } = primitive.outputs[0];
-		expect(connections).toHaveLength(0);
-
-		primitive.outputs[0].connectTo(n.inputs.ckpt_name);
-		({ connections } = primitive.outputs[0]);
-		expect(connections).toHaveLength(1);
-		expect(connections[0].targetNode.id).toBe(n.node.id);
-
-		// Convert back to widget and ensure input is removed
-		n.widgets.ckpt_name.convertToWidget();
-		expect(n.widgets.ckpt_name.isConvertedToInput).toBeFalsy();
-		expect(n.inputs.ckpt_name).toBeFalsy();
-		expect(n.inputs.length).toEqual(inputCount);
-	});
-
-	test("converted widget works on clone", async () => {
-		const { graph, ez } = await start();
-		let n = ez.CheckpointLoaderSimple();
-
-		// Convert the widget to an input
-		n.widgets.ckpt_name.convertToInput();
-		expect(n.widgets.ckpt_name.isConvertedToInput).toBeTruthy();
-
-		// Clone the node
-		n.menu["Clone"].call();
-		expect(graph.nodes).toHaveLength(2);
-		const clone = graph.nodes[1];
-		expect(clone.id).not.toEqual(n.id);
-
-		// Ensure the clone has an input
-		expect(clone.widgets.ckpt_name.isConvertedToInput).toBeTruthy();
-		expect(clone.inputs.ckpt_name).toBeTruthy();
-
-		// Ensure primitive connects to both nodes
-		let primitive = ez.PrimitiveNode();
-		primitive.outputs[0].connectTo(n.inputs.ckpt_name);
-		primitive.outputs[0].connectTo(clone.inputs.ckpt_name);
-		expect(primitive.outputs[0].connections).toHaveLength(2);
-
-		// Convert back to widget and ensure input is removed
-		clone.widgets.ckpt_name.convertToWidget();
-		expect(clone.widgets.ckpt_name.isConvertedToInput).toBeFalsy();
-		expect(clone.inputs.ckpt_name).toBeFalsy();
-	});
-
-	test("shows missing node error on custom node with converted input", async () => {
-		const { graph } = await start();
-
-		const dialogShow = jest.spyOn(graph.app.ui.dialog, "show");
-
-		await graph.app.loadGraphData({
-			last_node_id: 3,
-			last_link_id: 4,
-			nodes: [
-				{
-					id: 1,
-					type: "TestNode",
-					pos: [41.87329101561909, 389.7381480823742],
-					size: { 0: 220, 1: 374 },
-					flags: {},
-					order: 1,
-					mode: 0,
-					inputs: [{ name: "test", type: "FLOAT", link: 4, widget: { name: "test" }, slot_index: 0 }],
-					outputs: [],
-					properties: { "Node name for S&R": "TestNode" },
-					widgets_values: [1],
-				},
-				{
-					id: 3,
-					type: "PrimitiveNode",
-					pos: [-312, 433],
-					size: { 0: 210, 1: 82 },
-					flags: {},
-					order: 0,
-					mode: 0,
-					outputs: [{ links: [4], widget: { name: "test" } }],
-					title: "test",
-					properties: {},
-				},
-			],
-			links: [[4, 3, 0, 1, 6, "FLOAT"]],
-			groups: [],
-			config: {},
-			extra: {},
-			version: 0.4,
-		});
-
-		expect(dialogShow).toBeCalledTimes(1);
-		expect(dialogShow.mock.calls[0][0].innerHTML).toContain("the following node types were not found");
-		expect(dialogShow.mock.calls[0][0].innerHTML).toContain("TestNode");
-	});
-
-	test("defaultInput widgets can be converted back to inputs", async () => {
-		const { graph, ez } = await start({
-			mockNodeDefs: makeNodeDef("TestNode", { example: ["INT", { defaultInput: true }] }),
-		});
-
-		// Create test node and ensure it starts as an input
-		let n = ez.TestNode();
-		let w = n.widgets.example;
-		expect(w.isConvertedToInput).toBeTruthy();
-		let input = w.getConvertedInput();
-		expect(input).toBeTruthy();
-
-		// Ensure it can be converted to
-		w.convertToWidget();
-		expect(w.isConvertedToInput).toBeFalsy();
-		expect(n.inputs.length).toEqual(0);
-		// and from
-		w.convertToInput();
-		expect(w.isConvertedToInput).toBeTruthy();
-		input = w.getConvertedInput();
-
-		// Reload and ensure it still only has 1 converted widget
-		if (!assertNotNullOrUndefined(input)) return;
-
-		await connectPrimitiveAndReload(ez, graph, input, "number", 1);
-		n = graph.find(n);
-		expect(n.widgets).toHaveLength(1);
-		w = n.widgets.example;
-		expect(w.isConvertedToInput).toBeTruthy();
-
-		// Convert back to widget and ensure it is still a widget after reload
-		w.convertToWidget();
-		await graph.reload();
-		n = graph.find(n);
-		expect(n.widgets).toHaveLength(1);
-		expect(n.widgets[0].isConvertedToInput).toBeFalsy();
-		expect(n.inputs.length).toEqual(0);
-	});
-
-	test("forceInput widgets can not be converted back to inputs", async () => {
-		const { graph, ez } = await start({
-			mockNodeDefs: makeNodeDef("TestNode", { example: ["INT", { forceInput: true }] }),
-		});
-
-		// Create test node and ensure it starts as an input
-		let n = ez.TestNode();
-		let w = n.widgets.example;
-		expect(w.isConvertedToInput).toBeTruthy();
-		const input = w.getConvertedInput();
-		expect(input).toBeTruthy();
-
-		// Convert to widget should error
-		expect(() => w.convertToWidget()).toThrow();
-
-		// Reload and ensure it still only has 1 converted widget
-		if (assertNotNullOrUndefined(input)) {
-			await connectPrimitiveAndReload(ez, graph, input, "number", 1);
-			n = graph.find(n);
-			expect(n.widgets).toHaveLength(1);
-			expect(n.widgets.example.isConvertedToInput).toBeTruthy();
-		}
-	});
-
-	test("primitive can connect to matching combos on converted widgets", async () => {
-		const { ez } = await start({
-			mockNodeDefs: {
-				...makeNodeDef("TestNode1", { example: [["A", "B", "C"], { forceInput: true }] }),
-				...makeNodeDef("TestNode2", { example: [["A", "B", "C"], { forceInput: true }] }),
-			},
-		});
-
-		const n1 = ez.TestNode1();
-		const n2 = ez.TestNode2();
-		const p = ez.PrimitiveNode();
-		p.outputs[0].connectTo(n1.inputs[0]);
-		p.outputs[0].connectTo(n2.inputs[0]);
-		expect(p.outputs[0].connections).toHaveLength(2);
-		const valueWidget = p.widgets.value;
-		expect(valueWidget.widget.type).toBe("combo");
-		expect(valueWidget.widget.options.values).toEqual(["A", "B", "C"]);
-	});
-
-	test("primitive can not connect to non matching combos on converted widgets", async () => {
-		const { ez } = await start({
-			mockNodeDefs: {
-				...makeNodeDef("TestNode1", { example: [["A", "B", "C"], { forceInput: true }] }),
-				...makeNodeDef("TestNode2", { example: [["A", "B"], { forceInput: true }] }),
-			},
-		});
-
-		const n1 = ez.TestNode1();
-		const n2 = ez.TestNode2();
-		const p = ez.PrimitiveNode();
-		p.outputs[0].connectTo(n1.inputs[0]);
-		expect(() => p.outputs[0].connectTo(n2.inputs[0])).toThrow();
-		expect(p.outputs[0].connections).toHaveLength(1);
-	});
-
-	test("combo output can not connect to non matching combos list input", async () => {
-		const { ez } = await start({
-			mockNodeDefs: {
-				...makeNodeDef("TestNode1", {}, [["A", "B"]]),
-				...makeNodeDef("TestNode2", { example: [["A", "B"], { forceInput: true }] }),
-				...makeNodeDef("TestNode3", { example: [["A", "B", "C"], { forceInput: true }] }),
-			},
-		});
-
-		const n1 = ez.TestNode1();
-		const n2 = ez.TestNode2();
-		const n3 = ez.TestNode3();
-
-		n1.outputs[0].connectTo(n2.inputs[0]);
-		expect(() => n1.outputs[0].connectTo(n3.inputs[0])).toThrow();
-	});
-
-	test("combo primitive can filter list when control_after_generate called", async () => {
-		const { ez } = await start({
-			mockNodeDefs: {
-				...makeNodeDef("TestNode1", { example: [["A", "B", "C", "D", "AA", "BB", "CC", "DD", "AAA", "BBB"], {}] }),
-			},
-		});
-
-		const n1 = ez.TestNode1();
-		n1.widgets.example.convertToInput();
-		const p = ez.PrimitiveNode();
-		p.outputs[0].connectTo(n1.inputs[0]);
-
-		const value = p.widgets.value;
-		const control = p.widgets.control_after_generate.widget;
-		const filter = p.widgets.control_filter_list;
-
-		expect(p.widgets.length).toBe(3);
-		control.value = "increment";
-		expect(value.value).toBe("A");
-
-		// Manually trigger after queue when set to increment
-		control["afterQueued"]();
-		expect(value.value).toBe("B");
-
-		// Filter to items containing D
-		filter.value = "D";
-		control["afterQueued"]();
-		expect(value.value).toBe("D");
-		control["afterQueued"]();
-		expect(value.value).toBe("DD");
-
-		// Check decrement
-		value.value = "BBB";
-		control.value = "decrement";
-		filter.value = "B";
-		control["afterQueued"]();
-		expect(value.value).toBe("BB");
-		control["afterQueued"]();
-		expect(value.value).toBe("B");
-
-		// Check regex works
-		value.value = "BBB";
-		filter.value = "/[AB]|^C$/";
-		control["afterQueued"]();
-		expect(value.value).toBe("AAA");
-		control["afterQueued"]();
-		expect(value.value).toBe("BB");
-		control["afterQueued"]();
-		expect(value.value).toBe("AA");
-		control["afterQueued"]();
-		expect(value.value).toBe("C");
-		control["afterQueued"]();
-		expect(value.value).toBe("B");
-		control["afterQueued"]();
-		expect(value.value).toBe("A");
-
-		// Check random
-		control.value = "randomize";
-		filter.value = "/D/";
-		for (let i = 0; i < 100; i++) {
-			control["afterQueued"]();
-			expect(value.value === "D" || value.value === "DD").toBeTruthy();
-		}
-
-		// Ensure it doesnt apply when fixed
-		control.value = "fixed";
-		value.value = "B";
-		filter.value = "C";
-		control["afterQueued"]();
-		expect(value.value).toBe("B");
-	});
-
-	describe("reroutes", () => {
-		async function checkOutput(graph, values) {
-			expect((await graph.toPrompt()).output).toStrictEqual({
-				1: { inputs: { ckpt_name: "model1.safetensors" }, class_type: "CheckpointLoaderSimple" },
-				2: { inputs: { text: "positive", clip: ["1", 1] }, class_type: "CLIPTextEncode" },
-				3: { inputs: { text: "negative", clip: ["1", 1] }, class_type: "CLIPTextEncode" },
-				4: {
-					inputs: { width: values.width ?? 512, height: values.height ?? 512, batch_size: values?.batch_size ?? 1 },
-					class_type: "EmptyLatentImage",
-				},
-				5: {
-					inputs: {
-						seed: 0,
-						steps: 20,
-						cfg: 8,
-						sampler_name: "euler",
-						scheduler: values?.scheduler ?? "normal",
-						denoise: 1,
-						model: ["1", 0],
-						positive: ["2", 0],
-						negative: ["3", 0],
-						latent_image: ["4", 0],
-					},
-					class_type: "KSampler",
-				},
-				6: { inputs: { samples: ["5", 0], vae: ["1", 2] }, class_type: "VAEDecode" },
-				7: {
-					inputs: { filename_prefix: values.filename_prefix ?? "ComfyUI", images: ["6", 0] },
-					class_type: "SaveImage",
-				},
-			});
-		}
-
-		async function waitForWidget(node) {
-			// widgets are created slightly after the graph is ready
-			// hard to find an exact hook to get these so just wait for them to be ready
-			for (let i = 0; i < 10; i++) {
-				await new Promise((r) => setTimeout(r, 10));
-				if (node.widgets?.value) {
-					return;
-				}
-			}
-		}
-
-		it("can connect primitive via a reroute path to a widget input", async () => {
-			const { ez, graph } = await start();
-			const nodes = createDefaultWorkflow(ez, graph);
-
-			nodes.empty.widgets.width.convertToInput();
-			nodes.sampler.widgets.scheduler.convertToInput();
-			nodes.save.widgets.filename_prefix.convertToInput();
-
-			let widthReroute = ez.Reroute();
-			let schedulerReroute = ez.Reroute();
-			let fileReroute = ez.Reroute();
-
-			let widthNext = widthReroute;
-			let schedulerNext = schedulerReroute;
-			let fileNext = fileReroute;
-
-			for (let i = 0; i < 5; i++) {
-				let next = ez.Reroute();
-				widthNext.outputs[0].connectTo(next.inputs[0]);
-				widthNext = next;
-
-				next = ez.Reroute();
-				schedulerNext.outputs[0].connectTo(next.inputs[0]);
-				schedulerNext = next;
-
-				next = ez.Reroute();
-				fileNext.outputs[0].connectTo(next.inputs[0]);
-				fileNext = next;
-			}
-
-			widthNext.outputs[0].connectTo(nodes.empty.inputs.width);
-			schedulerNext.outputs[0].connectTo(nodes.sampler.inputs.scheduler);
-			fileNext.outputs[0].connectTo(nodes.save.inputs.filename_prefix);
-
-			let widthPrimitive = ez.PrimitiveNode();
-			let schedulerPrimitive = ez.PrimitiveNode();
-			let filePrimitive = ez.PrimitiveNode();
-
-			widthPrimitive.outputs[0].connectTo(widthReroute.inputs[0]);
-			schedulerPrimitive.outputs[0].connectTo(schedulerReroute.inputs[0]);
-			filePrimitive.outputs[0].connectTo(fileReroute.inputs[0]);
-			expect(widthPrimitive.widgets.value.value).toBe(512);
-			widthPrimitive.widgets.value.value = 1024;
-			expect(schedulerPrimitive.widgets.value.value).toBe("normal");
-			schedulerPrimitive.widgets.value.value = "simple";
-			expect(filePrimitive.widgets.value.value).toBe("ComfyUI");
-			filePrimitive.widgets.value.value = "ComfyTest";
-
-			await checkBeforeAndAfterReload(graph, async () => {
-				widthPrimitive = graph.find(widthPrimitive);
-				schedulerPrimitive = graph.find(schedulerPrimitive);
-				filePrimitive = graph.find(filePrimitive);
-				await waitForWidget(filePrimitive);
-				expect(widthPrimitive.widgets.length).toBe(2);
-				expect(schedulerPrimitive.widgets.length).toBe(3);
-				expect(filePrimitive.widgets.length).toBe(1);
-
-				await checkOutput(graph, {
-					width: 1024,
-					scheduler: "simple",
-					filename_prefix: "ComfyTest",
-				});
-			});
-		});
-		it("can connect primitive via a reroute path to multiple widget inputs", async () => {
-			const { ez, graph } = await start();
-			const nodes = createDefaultWorkflow(ez, graph);
-
-			nodes.empty.widgets.width.convertToInput();
-			nodes.empty.widgets.height.convertToInput();
-			nodes.empty.widgets.batch_size.convertToInput();
-
-			let reroute = ez.Reroute();
-			let prevReroute = reroute;
-			for (let i = 0; i < 5; i++) {
-				const next = ez.Reroute();
-				prevReroute.outputs[0].connectTo(next.inputs[0]);
-				prevReroute = next;
-			}
-
-			const r1 = ez.Reroute(prevReroute.outputs[0]);
-			const r2 = ez.Reroute(prevReroute.outputs[0]);
-			const r3 = ez.Reroute(r2.outputs[0]);
-			const r4 = ez.Reroute(r2.outputs[0]);
-
-			r1.outputs[0].connectTo(nodes.empty.inputs.width);
-			r3.outputs[0].connectTo(nodes.empty.inputs.height);
-			r4.outputs[0].connectTo(nodes.empty.inputs.batch_size);
-
-			let primitive = ez.PrimitiveNode();
-			primitive.outputs[0].connectTo(reroute.inputs[0]);
-			expect(primitive.widgets.value.value).toBe(1);
-			primitive.widgets.value.value = 64;
-
-			await checkBeforeAndAfterReload(graph, async (r) => {
-				primitive = graph.find(primitive);
-				await waitForWidget(primitive);
-
-				// Ensure widget configs are merged
-				expect(primitive.widgets.value.widget.options?.min).toBe(16); // width/height min
-				expect(primitive.widgets.value.widget.options?.max).toBe(4096); // batch max
-				expect(primitive.widgets.value.widget.options?.step).toBe(80); // width/height step * 10
-
-				await checkOutput(graph, {
-					width: 64,
-					height: 64,
-					batch_size: 64,
-				});
-			});
-		});
-	});
-});

tests-ui/utils/ezgraph.js

@@ -1,452 +0,0 @@
-// @ts-check
-/// <reference path="../../web/types/litegraph.d.ts" />
-
-/**
- * @typedef { import("../../web/scripts/app")["app"] } app
- * @typedef { import("../../web/types/litegraph") } LG
- * @typedef { import("../../web/types/litegraph").IWidget } IWidget
- * @typedef { import("../../web/types/litegraph").ContextMenuItem } ContextMenuItem
- * @typedef { import("../../web/types/litegraph").INodeInputSlot } INodeInputSlot
- * @typedef { import("../../web/types/litegraph").INodeOutputSlot } INodeOutputSlot
- * @typedef { InstanceType<LG["LGraphNode"]> & { widgets?: Array<IWidget> } } LGNode
- * @typedef { (...args: EzOutput[] | [...EzOutput[], Record<string, unknown>]) => EzNode } EzNodeFactory
- */
-
-export class EzConnection {
-	/** @type { app } */
-	app;
-	/** @type { InstanceType<LG["LLink"]> } */
-	link;
-
-	get originNode() {
-		return new EzNode(this.app, this.app.graph.getNodeById(this.link.origin_id));
-	}
-
-	get originOutput() {
-		return this.originNode.outputs[this.link.origin_slot];
-	}
-
-	get targetNode() {
-		return new EzNode(this.app, this.app.graph.getNodeById(this.link.target_id));
-	}
-
-	get targetInput() {
-		return this.targetNode.inputs[this.link.target_slot];
-	}
-
-	/**
-	 * @param { app } app
-	 * @param { InstanceType<LG["LLink"]> } link
-	 */
-	constructor(app, link) {
-		this.app = app;
-		this.link = link;
-	}
-
-	disconnect() {
-		this.targetInput.disconnect();
-	}
-}
-
-export class EzSlot {
-	/** @type { EzNode } */
-	node;
-	/** @type { number } */
-	index;
-
-	/**
-	 * @param { EzNode } node
-	 * @param { number } index
-	 */
-	constructor(node, index) {
-		this.node = node;
-		this.index = index;
-	}
-}
-
-export class EzInput extends EzSlot {
-	/** @type { INodeInputSlot } */
-	input;
-
-	/**
-	 * @param { EzNode } node
-	 * @param { number } index
-	 * @param { INodeInputSlot } input
-	 */
-	constructor(node, index, input) {
-		super(node, index);
-		this.input = input;
-	}
-
-	get connection() {
-		const link = this.node.node.inputs?.[this.index]?.link;
-		if (link == null) {
-			return null;
-		}
-		return new EzConnection(this.node.app, this.node.app.graph.links[link]);
-	}
-
-	disconnect() {
-		this.node.node.disconnectInput(this.index);
-	}
-}
-
-export class EzOutput extends EzSlot {
-	/** @type { INodeOutputSlot } */
-	output;
-
-	/**
-	 * @param { EzNode } node
-	 * @param { number } index
-	 * @param { INodeOutputSlot } output
-	 */
-	constructor(node, index, output) {
-		super(node, index);
-		this.output = output;
-	}
-
-	get connections() {
-		return (this.node.node.outputs?.[this.index]?.links ?? []).map(
-			(l) => new EzConnection(this.node.app, this.node.app.graph.links[l])
-		);
-	}
-
-	/**
-	 * @param { EzInput } input
-	 */
-	connectTo(input) {
-		if (!input) throw new Error("Invalid input");
-
-		/**
-		 * @type { LG["LLink"] | null }
-		 */
-		const link = this.node.node.connect(this.index, input.node.node, input.index);
-		if (!link) {
-			const inp = input.input;
-			const inName = inp.name || inp.label || inp.type;
-			throw new Error(
-				`Connecting from ${input.node.node.type}#${input.node.id}[${inName}#${input.index}] -> ${this.node.node.type}#${this.node.id}[${
-					this.output.name ?? this.output.type
-				}#${this.index}] failed.`
-			);
-		}
-		return link;
-	}
-}
-
-export class EzNodeMenuItem {
-	/** @type { EzNode } */
-	node;
-	/** @type { number } */
-	index;
-	/** @type { ContextMenuItem } */
-	item;
-
-	/**
-	 * @param { EzNode } node
-	 * @param { number } index
-	 * @param { ContextMenuItem } item
-	 */
-	constructor(node, index, item) {
-		this.node = node;
-		this.index = index;
-		this.item = item;
-	}
-
-	call(selectNode = true) {
-		if (!this.item?.callback) throw new Error(`Menu Item ${this.item?.content ?? "[null]"} has no callback.`);
-		if (selectNode) {
-			this.node.select();
-		}
-		return this.item.callback.call(this.node.node, undefined, undefined, undefined, undefined, this.node.node);
-	}
-}
-
-export class EzWidget {
-	/** @type { EzNode } */
-	node;
-	/** @type { number } */
-	index;
-	/** @type { IWidget } */
-	widget;
-
-	/**
-	 * @param { EzNode } node
-	 * @param { number } index
-	 * @param { IWidget } widget
-	 */
-	constructor(node, index, widget) {
-		this.node = node;
-		this.index = index;
-		this.widget = widget;
-	}
-
-	get value() {
-		return this.widget.value;
-	}
-
-	set value(v) {
-		this.widget.value = v;
-		this.widget.callback?.call?.(this.widget, v)
-	}
-
-	get isConvertedToInput() {
-		// @ts-ignore : this type is valid for converted widgets
-		return this.widget.type === "converted-widget";
-	}
-
-	getConvertedInput() {
-		if (!this.isConvertedToInput) throw new Error(`Widget ${this.widget.name} is not converted to input.`);
-
-		return this.node.inputs.find((inp) => inp.input["widget"]?.name === this.widget.name);
-	}
-
-	convertToWidget() {
-		if (!this.isConvertedToInput)
-			throw new Error(`Widget ${this.widget.name} cannot be converted as it is already a widget.`);
-		var menu = this.node.menu["Convert Input to Widget"].item.submenu.options;
-		var index = menu.findIndex(a => a.content == `Convert ${this.widget.name} to widget`);
-		menu[index].callback.call();
-	}
-
-	convertToInput() {
-		if (this.isConvertedToInput)
-			throw new Error(`Widget ${this.widget.name} cannot be converted as it is already an input.`);
-		var menu = this.node.menu["Convert Widget to Input"].item.submenu.options;
-		var index = menu.findIndex(a => a.content == `Convert ${this.widget.name} to input`);
-		menu[index].callback.call();
-	}
-}
-
-export class EzNode {
-	/** @type { app } */
-	app;
-	/** @type { LGNode } */
-	node;
-
-	/**
-	 * @param { app } app
-	 * @param { LGNode } node
-	 */
-	constructor(app, node) {
-		this.app = app;
-		this.node = node;
-	}
-
-	get id() {
-		return this.node.id;
-	}
-
-	get inputs() {
-		return this.#makeLookupArray("inputs", "name", EzInput);
-	}
-
-	get outputs() {
-		return this.#makeLookupArray("outputs", "name", EzOutput);
-	}
-
-	get widgets() {
-		return this.#makeLookupArray("widgets", "name", EzWidget);
-	}
-
-	get menu() {
-		return this.#makeLookupArray(() => this.app.canvas.getNodeMenuOptions(this.node), "content", EzNodeMenuItem);
-	}
-
-	get isRemoved() {
-		return !this.app.graph.getNodeById(this.id);
-	}
-
-	select(addToSelection = false) {
-		this.app.canvas.selectNode(this.node, addToSelection);
-	}
-
-	// /**
-	//  * @template { "inputs" | "outputs" } T
-	//  * @param { T } type
-	//  * @returns { Record<string, type extends "inputs" ? EzInput : EzOutput> & (type extends "inputs" ? EzInput [] : EzOutput[]) }
-	//  */
-	// #getSlotItems(type) {
-	// 	// @ts-ignore : these items are correct
-	// 	return (this.node[type] ?? []).reduce((p, s, i) => {
-	// 		if (s.name in p) {
-	// 			throw new Error(`Unable to store input ${s.name} on array as name conflicts.`);
-	// 		}
-	// 		// @ts-ignore
-	// 		p.push((p[s.name] = new (type === "inputs" ? EzInput : EzOutput)(this, i, s)));
-	// 		return p;
-	// 	}, Object.assign([], { $: this }));
-	// }
-
-	/**
-	 * @template { { new(node: EzNode, index: number, obj: any): any } } T
-	 * @param { "inputs" | "outputs" | "widgets" | (() => Array<unknown>) } nodeProperty
-	 * @param { string } nameProperty
-	 * @param { T } ctor
-	 * @returns { Record<string, InstanceType<T>> & Array<InstanceType<T>> }
-	 */
-	#makeLookupArray(nodeProperty, nameProperty, ctor) {
-		const items = typeof nodeProperty === "function" ? nodeProperty() : this.node[nodeProperty];
-		// @ts-ignore
-		return (items ?? []).reduce((p, s, i) => {
-			if (!s) return p;
-
-			const name = s[nameProperty];
-			const item = new ctor(this, i, s);
-			// @ts-ignore
-			p.push(item);
-			if (name) {
-				// @ts-ignore
-				if (name in p) {
-					throw new Error(`Unable to store ${nodeProperty} ${name} on array as name conflicts.`);
-				}
-			}
-			// @ts-ignore
-			p[name] = item;
-			return p;
-		}, Object.assign([], { $: this }));
-	}
-}
-
-export class EzGraph {
-	/** @type { app } */
-	app;
-
-	/**
-	 * @param { app } app
-	 */
-	constructor(app) {
-		this.app = app;
-	}
-
-	get nodes() {
-		return this.app.graph._nodes.map((n) => new EzNode(this.app, n));
-	}
-
-	clear() {
-		this.app.graph.clear();
-	}
-
-	arrange() {
-		this.app.graph.arrange();
-	}
-
-	stringify() {
-		return JSON.stringify(this.app.graph.serialize(), undefined);
-	}
-
-	/**
-	 * @param { number | LGNode | EzNode } obj
-	 * @returns { EzNode }
-	 */
-	find(obj) {
-		let match;
-		let id;
-		if (typeof obj === "number") {
-			id = obj;
-		} else {
-			id = obj.id;
-		}
-
-		match = this.app.graph.getNodeById(id);
-
-		if (!match) {
-			throw new Error(`Unable to find node with ID ${id}.`);
-		}
-
-		return new EzNode(this.app, match);
-	}
-
-	/**
-	 * @returns { Promise<void> }
-	 */
-	reload() {
-		const graph = JSON.parse(JSON.stringify(this.app.graph.serialize()));
-		return new Promise((r) => {
-			this.app.graph.clear();
-			setTimeout(async () => {
-				await this.app.loadGraphData(graph);
-				r();
-			}, 10);
-		});
-	}
-
-	/**
-	 * @returns { Promise<{
-	 * 	workflow: {},
-	 * 	output: Record<string, {
-	 * 		class_name: string,
-	 * 		inputs: Record<string, [string, number] | unknown>
-	 * }>}> }
-	 */
-	toPrompt() {
-		// @ts-ignore
-		return this.app.graphToPrompt();
-	}
-}
-
-export const Ez = {
-	/**
-	 * Quickly build and interact with a ComfyUI graph
-	 * @example
-	 * const { ez, graph } = Ez.graph(app);
-	 * graph.clear();
-	 * const [model, clip, vae] = ez.CheckpointLoaderSimple().outputs;
-	 * const [pos] = ez.CLIPTextEncode(clip, { text: "positive" }).outputs;
-	 * const [neg] = ez.CLIPTextEncode(clip, { text: "negative" }).outputs;
-	 * const [latent] = ez.KSampler(model, pos, neg, ...ez.EmptyLatentImage().outputs).outputs;
-	 * const [image] = ez.VAEDecode(latent, vae).outputs;
-	 * const saveNode = ez.SaveImage(image);
-	 * console.log(saveNode);
-	 * graph.arrange();
-	 * @param { app } app
-	 * @param { LG["LiteGraph"] } LiteGraph
-	 * @param { LG["LGraphCanvas"] } LGraphCanvas
-	 * @param { boolean } clearGraph
-	 * @returns { { graph: EzGraph, ez: Record<string, EzNodeFactory> } }
-	 */
-	graph(app, LiteGraph = window["LiteGraph"], LGraphCanvas = window["LGraphCanvas"], clearGraph = true) {
-		// Always set the active canvas so things work
-		LGraphCanvas.active_canvas = app.canvas;
-
-		if (clearGraph) {
-			app.graph.clear();
-		}
-
-		// @ts-ignore : this proxy handles utility methods & node creation
-		const factory = new Proxy(
-			{},
-			{
-				get(_, p) {
-					if (typeof p !== "string") throw new Error("Invalid node");
-					const node = LiteGraph.createNode(p);
-					if (!node) throw new Error(`Unknown node "${p}"`);
-					app.graph.add(node);
-
-					/**
-					 * @param {Parameters<EzNodeFactory>} args
-					 */
-					return function (...args) {
-						const ezNode = new EzNode(app, node);
-						const inputs = ezNode.inputs;
-
-						let slot = 0;
-						for (const arg of args) {
-							if (arg instanceof EzOutput) {
-								arg.connectTo(inputs[slot++]);
-							} else {
-								for (const k in arg) {
-									ezNode.widgets[k].value = arg[k];
-								}
-							}
-						}
-
-						return ezNode;
-					};
-				},
-			}
-		);
-
-		return { graph: new EzGraph(app), ez: factory };
-	},
-};

tests-ui/utils/index.js

@@ -1,129 +0,0 @@
-const { mockApi } = require("./setup");
-const { Ez } = require("./ezgraph");
-const lg = require("./litegraph");
-const fs = require("fs");
-const path = require("path");
-
-const html = fs.readFileSync(path.resolve(__dirname, "../../web/index.html"))
-
-/**
- *
- * @param { Parameters<typeof mockApi>[0] & { 
- * 	resetEnv?: boolean, 
- * 	preSetup?(app): Promise<void>,
- *  localStorage?: Record<string, string> 
- * } } config
- * @returns
- */
-export async function start(config = {}) {
-	if(config.resetEnv) {
-		jest.resetModules();
-		jest.resetAllMocks();
-        lg.setup(global);
-		localStorage.clear();
-		sessionStorage.clear();
-	}
-
-	Object.assign(localStorage, config.localStorage ?? {});
-	document.body.innerHTML = html;
-
-	mockApi(config);
-	const { app } = require("../../web/scripts/app");
-	config.preSetup?.(app);
-	await app.setup();
-
-	return { ...Ez.graph(app, global["LiteGraph"], global["LGraphCanvas"]), app };
-}
-
-/**
- * @param { ReturnType<Ez["graph"]>["graph"] } graph
- * @param { (hasReloaded: boolean) => (Promise<void> | void) } cb
- */
-export async function checkBeforeAndAfterReload(graph, cb) {
-	await cb(false);
-	await graph.reload();
-	await cb(true);
-}
-
-/**
- * @param { string } name
- * @param { Record<string, string | [string | string[], any]> } input
- * @param { (string | string[])[] | Record<string, string | string[]> } output
- * @returns { Record<string, import("../../web/types/comfy").ComfyObjectInfo> }
- */
-export function makeNodeDef(name, input, output = {}) {
-	const nodeDef = {
-		name,
-		category: "test",
-		output: [],
-		output_name: [],
-		output_is_list: [],
-		input: {
-			required: {},
-		},
-	};
-	for (const k in input) {
-		nodeDef.input.required[k] = typeof input[k] === "string" ? [input[k], {}] : [...input[k]];
-	}
-	if (output instanceof Array) {
-		output = output.reduce((p, c) => {
-			p[c] = c;
-			return p;
-		}, {});
-	}
-	for (const k in output) {
-		nodeDef.output.push(output[k]);
-		nodeDef.output_name.push(k);
-		nodeDef.output_is_list.push(false);
-	}
-
-	return { [name]: nodeDef };
-}
-
-/**
-/**
- * @template { any } T
- * @param { T } x
- * @returns { x is Exclude<T, null | undefined> }
- */
-export function assertNotNullOrUndefined(x) {
-	expect(x).not.toEqual(null);
-	expect(x).not.toEqual(undefined);
-	return true;
-}
-
-/**
- *
- * @param { ReturnType<Ez["graph"]>["ez"] } ez
- * @param { ReturnType<Ez["graph"]>["graph"] } graph
- */
-export function createDefaultWorkflow(ez, graph) {
-	graph.clear();
-	const ckpt = ez.CheckpointLoaderSimple();
-
-	const pos = ez.CLIPTextEncode(ckpt.outputs.CLIP, { text: "positive" });
-	const neg = ez.CLIPTextEncode(ckpt.outputs.CLIP, { text: "negative" });
-
-	const empty = ez.EmptyLatentImage();
-	const sampler = ez.KSampler(
-		ckpt.outputs.MODEL,
-		pos.outputs.CONDITIONING,
-		neg.outputs.CONDITIONING,
-		empty.outputs.LATENT
-	);
-
-	const decode = ez.VAEDecode(sampler.outputs.LATENT, ckpt.outputs.VAE);
-	const save = ez.SaveImage(decode.outputs.IMAGE);
-	graph.arrange();
-
-	return { ckpt, pos, neg, empty, sampler, decode, save };
-}
-
-export async function getNodeDefs() {
-	const { api } = require("../../web/scripts/api");
-	return api.getNodeDefs();
-}
-
-export async function getNodeDef(nodeId) {
-	return (await getNodeDefs())[nodeId];
-}

tests-ui/utils/litegraph.js

@@ -1,36 +0,0 @@
-const fs = require("fs");
-const path = require("path");
-const { nop } = require("../utils/nopProxy");
-
-function forEachKey(cb) {
-	for (const k of [
-		"LiteGraph",
-		"LGraph",
-		"LLink",
-		"LGraphNode",
-		"LGraphGroup",
-		"DragAndScale",
-		"LGraphCanvas",
-		"ContextMenu",
-	]) {
-		cb(k);
-	}
-}
-
-export function setup(ctx) {
-	const lg = fs.readFileSync(path.resolve("../web/lib/litegraph.core.js"), "utf-8");
-	const globalTemp = {};
-	(function (console) {
-		eval(lg);
-	}).call(globalTemp, nop);
-
-	forEachKey((k) => (ctx[k] = globalTemp[k]));
-	require(path.resolve("../web/lib/litegraph.extensions.js"));
-}
-
-export function teardown(ctx) {
-	forEachKey((k) => delete ctx[k]);
-
-	// Clear document after each run
-	document.getElementsByTagName("html")[0].innerHTML = ""; 
-}

tests-ui/utils/nopProxy.js

@@ -1,6 +0,0 @@
-export const nop = new Proxy(function () {}, {
-	get: () => nop,
-	set: () => true,
-	apply: () => nop,
-	construct: () => nop,
-});

tests-ui/utils/setup.js

@@ -1,82 +0,0 @@
-require("../../web/scripts/api");
-
-const fs = require("fs");
-const path = require("path");
-function* walkSync(dir) {
-	const files = fs.readdirSync(dir, { withFileTypes: true });
-	for (const file of files) {
-		if (file.isDirectory()) {
-			yield* walkSync(path.join(dir, file.name));
-		} else {
-			yield path.join(dir, file.name);
-		}
-	}
-}
-
-/**
- * @typedef { import("../../web/types/comfy").ComfyObjectInfo } ComfyObjectInfo
- */
-
-/**
- * @param {{ 
- * 	mockExtensions?: string[], 
- * 	mockNodeDefs?: Record<string, ComfyObjectInfo>,
-* 	settings?: Record<string, string>
-* 	userConfig?: {storage: "server" | "browser", users?: Record<string, any>, migrated?: boolean },
-* 	userData?: Record<string, any>
- * }} config
- */
-export function mockApi(config = {}) {
-	let { mockExtensions, mockNodeDefs, userConfig, settings, userData } = {
-		userConfig,
-		settings: {},
-		userData: {},
-		...config,
-	};
-	if (!mockExtensions) {
-		mockExtensions = Array.from(walkSync(path.resolve("../web/extensions/core")))
-			.filter((x) => x.endsWith(".js"))
-			.map((x) => path.relative(path.resolve("../web"), x));
-	}
-	if (!mockNodeDefs) {
-		mockNodeDefs = JSON.parse(fs.readFileSync(path.resolve("./data/object_info.json")));
-	}
-
-	const events = new EventTarget();
-	const mockApi = {
-		addEventListener: events.addEventListener.bind(events),
-		removeEventListener: events.removeEventListener.bind(events),
-		dispatchEvent: events.dispatchEvent.bind(events),
-		getSystemStats: jest.fn(),
-		getExtensions: jest.fn(() => mockExtensions),
-		getNodeDefs: jest.fn(() => mockNodeDefs),
-		init: jest.fn(),
-		apiURL: jest.fn((x) => "../../web/" + x),
-		createUser: jest.fn((username) => {
-			if(username in userConfig.users) {
-				return { status: 400, json: () => "Duplicate" }
-			}
-			userConfig.users[username + "!"] = username;
-			return { status: 200, json: () => username + "!" }
-		}),
-		getUserConfig: jest.fn(() => userConfig ?? { storage: "browser", migrated: false }),
-		getSettings: jest.fn(() => settings),
-		storeSettings: jest.fn((v) => Object.assign(settings, v)),
-		getUserData: jest.fn((f) => {
-			if (f in userData) {
-				return { status: 200, json: () => userData[f] };
-			} else {
-				return { status: 404 };
-			}
-		}),
-		storeUserData: jest.fn((file, data) => {
-			userData[file] = data;
-		}),
-		listUserData: jest.fn(() => [])
-	};
-	jest.mock("../../web/scripts/api", () => ({
-		get api() {
-			return mockApi;
-		},
-	}));
-}

tests-unit/prompt_server_test/download_models_test.py

@@ -0,0 +1,321 @@
+import pytest
+import aiohttp
+from aiohttp import ClientResponse
+import itertools
+import os 
+from unittest.mock import AsyncMock, patch, MagicMock
+from model_filemanager import download_model, validate_model_subdirectory, track_download_progress, create_model_path, check_file_exists, DownloadStatusType, DownloadModelStatus, validate_filename
+
+class AsyncIteratorMock:
+    """
+    A mock class that simulates an asynchronous iterator.
+    This is used to mimic the behavior of aiohttp's content iterator.
+    """
+    def __init__(self, seq):
+        # Convert the input sequence into an iterator
+        self.iter = iter(seq)
+
+    def __aiter__(self):
+        # This method is called when 'async for' is used
+        return self
+
+    async def __anext__(self):
+        # This method is called for each iteration in an 'async for' loop
+        try:
+            return next(self.iter)
+        except StopIteration:
+            # This is the asynchronous equivalent of StopIteration
+            raise StopAsyncIteration
+
+class ContentMock:
+    """
+    A mock class that simulates the content attribute of an aiohttp ClientResponse.
+    This class provides the iter_chunked method which returns an async iterator of chunks.
+    """
+    def __init__(self, chunks):
+        # Store the chunks that will be returned by the iterator
+        self.chunks = chunks
+
+    def iter_chunked(self, chunk_size):
+        # This method mimics aiohttp's content.iter_chunked()
+        # For simplicity in testing, we ignore chunk_size and just return our predefined chunks
+        return AsyncIteratorMock(self.chunks)
+
+@pytest.mark.asyncio
+async def test_download_model_success():
+    mock_response = AsyncMock(spec=aiohttp.ClientResponse)
+    mock_response.status = 200
+    mock_response.headers = {'Content-Length': '1000'}
+    # Create a mock for content that returns an async iterator directly
+    chunks = [b'a' * 500, b'b' * 300, b'c' * 200]
+    mock_response.content = ContentMock(chunks)
+
+    mock_make_request = AsyncMock(return_value=mock_response)
+    mock_progress_callback = AsyncMock()
+
+    # Mock file operations
+    mock_open = MagicMock()
+    mock_file = MagicMock()
+    mock_open.return_value.__enter__.return_value = mock_file
+    time_values = itertools.count(0, 0.1)
+
+    with patch('model_filemanager.create_model_path', return_value=('models/checkpoints/model.sft', 'checkpoints/model.sft')), \
+         patch('model_filemanager.check_file_exists', return_value=None), \
+         patch('builtins.open', mock_open), \
+         patch('time.time', side_effect=time_values):  # Simulate time passing
+
+        result = await download_model(
+            mock_make_request,
+            'model.sft',
+            'http://example.com/model.sft',
+            'checkpoints',
+            mock_progress_callback
+        )
+
+    # Assert the result
+    assert isinstance(result, DownloadModelStatus)
+    assert result.message == 'Successfully downloaded model.sft'
+    assert result.status == 'completed'
+    assert result.already_existed is False
+
+    # Check progress callback calls
+    assert mock_progress_callback.call_count >= 3  # At least start, one progress update, and completion
+    
+    # Check initial call
+    mock_progress_callback.assert_any_call(
+        'checkpoints/model.sft',
+        DownloadModelStatus(DownloadStatusType.PENDING, 0, "Starting download of model.sft", False)
+    )
+
+    # Check final call
+    mock_progress_callback.assert_any_call(
+        'checkpoints/model.sft',
+        DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "Successfully downloaded model.sft", False)
+    )
+
+    # Verify file writing
+    mock_file.write.assert_any_call(b'a' * 500)
+    mock_file.write.assert_any_call(b'b' * 300)
+    mock_file.write.assert_any_call(b'c' * 200)
+
+    # Verify request was made
+    mock_make_request.assert_called_once_with('http://example.com/model.sft')
+
+@pytest.mark.asyncio
+async def test_download_model_url_request_failure():
+    # Mock dependencies
+    mock_response = AsyncMock(spec=ClientResponse)
+    mock_response.status = 404  # Simulate a "Not Found" error
+    mock_get = AsyncMock(return_value=mock_response)
+    mock_progress_callback = AsyncMock()
+
+    # Mock the create_model_path function
+    with patch('model_filemanager.create_model_path', return_value=('/mock/path/model.safetensors', 'mock/path/model.safetensors')):
+        # Mock the check_file_exists function to return None (file doesn't exist)
+        with patch('model_filemanager.check_file_exists', return_value=None):
+            # Call the function
+            result = await download_model(
+                mock_get,
+                'model.safetensors',
+                'http://example.com/model.safetensors',
+                'mock_directory',
+                mock_progress_callback
+            )
+
+    # Assert the expected behavior
+    assert isinstance(result, DownloadModelStatus)
+    assert result.status == 'error'
+    assert result.message == 'Failed to download model.safetensors. Status code: 404'
+    assert result.already_existed is False
+
+    # Check that progress_callback was called with the correct arguments
+    mock_progress_callback.assert_any_call(
+        'mock_directory/model.safetensors',
+        DownloadModelStatus(
+            status=DownloadStatusType.PENDING,
+            progress_percentage=0,
+            message='Starting download of model.safetensors',
+            already_existed=False
+        )
+    )
+    mock_progress_callback.assert_called_with(
+        'mock_directory/model.safetensors',
+        DownloadModelStatus(
+            status=DownloadStatusType.ERROR,
+            progress_percentage=0,
+            message='Failed to download model.safetensors. Status code: 404',
+            already_existed=False
+        )
+    )
+
+    # Verify that the get method was called with the correct URL
+    mock_get.assert_called_once_with('http://example.com/model.safetensors')
+
+@pytest.mark.asyncio
+async def test_download_model_invalid_model_subdirectory():
+    
+    mock_make_request = AsyncMock()
+    mock_progress_callback = AsyncMock()
+
+    
+    result = await download_model(
+        mock_make_request,
+        'model.sft',
+        'http://example.com/model.sft',
+        '../bad_path',
+        mock_progress_callback
+    )
+
+    # Assert the result
+    assert isinstance(result, DownloadModelStatus)
+    assert result.message == 'Invalid model subdirectory'
+    assert result.status == 'error'
+    assert result.already_existed is False
+
+
+# For create_model_path function
+def test_create_model_path(tmp_path, monkeypatch):
+    mock_models_dir = tmp_path / "models"
+    monkeypatch.setattr('folder_paths.models_dir', str(mock_models_dir))
+    
+    model_name = "test_model.sft"
+    model_directory = "test_dir"
+    
+    file_path, relative_path = create_model_path(model_name, model_directory, mock_models_dir)
+    
+    assert file_path == str(mock_models_dir / model_directory / model_name)
+    assert relative_path == f"{model_directory}/{model_name}"
+    assert os.path.exists(os.path.dirname(file_path))
+
+
+@pytest.mark.asyncio
+async def test_check_file_exists_when_file_exists(tmp_path):
+    file_path = tmp_path / "existing_model.sft"
+    file_path.touch()  # Create an empty file
+    
+    mock_callback = AsyncMock()
+    
+    result = await check_file_exists(str(file_path), "existing_model.sft", mock_callback, "test/existing_model.sft")
+    
+    assert result is not None
+    assert result.status == "completed"
+    assert result.message == "existing_model.sft already exists"
+    assert result.already_existed is True
+    
+    mock_callback.assert_called_once_with(
+        "test/existing_model.sft",
+        DownloadModelStatus(DownloadStatusType.COMPLETED, 100, "existing_model.sft already exists", already_existed=True)
+    )
+
+@pytest.mark.asyncio
+async def test_check_file_exists_when_file_does_not_exist(tmp_path):
+    file_path = tmp_path / "non_existing_model.sft"
+    
+    mock_callback = AsyncMock()
+    
+    result = await check_file_exists(str(file_path), "non_existing_model.sft", mock_callback, "test/non_existing_model.sft")
+    
+    assert result is None
+    mock_callback.assert_not_called()
+
+@pytest.mark.asyncio
+async def test_track_download_progress_no_content_length():
+    mock_response = AsyncMock(spec=aiohttp.ClientResponse)
+    mock_response.headers = {}  # No Content-Length header
+    mock_response.content.iter_chunked.return_value = AsyncIteratorMock([b'a' * 500, b'b' * 500])
+
+    mock_callback = AsyncMock()
+    mock_open = MagicMock(return_value=MagicMock())
+
+    with patch('builtins.open', mock_open):
+        result = await track_download_progress(
+            mock_response, '/mock/path/model.sft', 'model.sft',
+            mock_callback, 'models/model.sft', interval=0.1
+        )
+
+    assert result.status == "completed"
+    # Check that progress was reported even without knowing the total size
+    mock_callback.assert_any_call(
+        'models/model.sft',
+        DownloadModelStatus(DownloadStatusType.IN_PROGRESS, 0, "Downloading model.sft", already_existed=False)
+    )
+
+@pytest.mark.asyncio
+async def test_track_download_progress_interval():
+    mock_response = AsyncMock(spec=aiohttp.ClientResponse)
+    mock_response.headers = {'Content-Length': '1000'}
+    mock_response.content.iter_chunked.return_value = AsyncIteratorMock([b'a' * 100] * 10)
+
+    mock_callback = AsyncMock()
+    mock_open = MagicMock(return_value=MagicMock())
+
+    # Create a mock time function that returns incremental float values
+    mock_time = MagicMock()
+    mock_time.side_effect = [i * 0.5 for i in range(30)]  # This should be enough for 10 chunks
+
+    with patch('builtins.open', mock_open), \
+         patch('time.time', mock_time):
+        await track_download_progress(
+            mock_response, '/mock/path/model.sft', 'model.sft',
+            mock_callback, 'models/model.sft', interval=1.0
+        )
+
+    # Print out the actual call count and the arguments of each call for debugging
+    print(f"mock_callback was called {mock_callback.call_count} times")
+    for i, call in enumerate(mock_callback.call_args_list):
+        args, kwargs = call
+        print(f"Call {i + 1}: {args[1].status}, Progress: {args[1].progress_percentage:.2f}%")
+
+    # Assert that progress was updated at least 3 times (start, at least one interval, and end)
+    assert mock_callback.call_count >= 3, f"Expected at least 3 calls, but got {mock_callback.call_count}"
+
+    # Verify the first and last calls
+    first_call = mock_callback.call_args_list[0]
+    assert first_call[0][1].status == "in_progress"
+    # Allow for some initial progress, but it should be less than 50%
+    assert 0 <= first_call[0][1].progress_percentage < 50, f"First call progress was {first_call[0][1].progress_percentage}%"
+
+    last_call = mock_callback.call_args_list[-1]
+    assert last_call[0][1].status == "completed"
+    assert last_call[0][1].progress_percentage == 100
+
+def test_valid_subdirectory():
+    assert validate_model_subdirectory("valid-model123") is True
+
+def test_subdirectory_too_long():
+    assert validate_model_subdirectory("a" * 51) is False
+
+def test_subdirectory_with_double_dots():
+    assert validate_model_subdirectory("model/../unsafe") is False
+
+def test_subdirectory_with_slash():
+    assert validate_model_subdirectory("model/unsafe") is False
+
+def test_subdirectory_with_special_characters():
+    assert validate_model_subdirectory("model@unsafe") is False
+
+def test_subdirectory_with_underscore_and_dash():
+    assert validate_model_subdirectory("valid_model-name") is True
+
+def test_empty_subdirectory():
+    assert validate_model_subdirectory("") is False
+
+@pytest.mark.parametrize("filename, expected", [
+    ("valid_model.safetensors", True),
+    ("valid_model.sft", True),
+    ("valid model.safetensors", True), # Test with space
+    ("UPPERCASE_MODEL.SAFETENSORS", True),
+    ("model_with.multiple.dots.pt", False),
+    ("", False),  # Empty string
+    ("../../../etc/passwd", False),  # Path traversal attempt
+    ("/etc/passwd", False),  # Absolute path
+    ("\\windows\\system32\\config\\sam", False),  # Windows path
+    (".hidden_file.pt", False),  # Hidden file
+    ("invalid<char>.ckpt", False),  # Invalid character
+    ("invalid?.ckpt", False),  # Another invalid character
+    ("very" * 100 + ".safetensors", False),  # Too long filename
+    ("\nmodel_with_newline.pt", False),  # Newline character
+    ("model_with_emoji😊.pt", False),  # Emoji in filename
+])
+def test_validate_filename(filename, expected):
+    assert validate_filename(filename) == expected

tests-unit/requirements.txt

@@ -1 +1,3 @@
 pytest>=7.8.0
+pytest-aiohttp
+pytest-asyncio

tests/inference/extra_model_paths.yaml

@@ -0,0 +1,4 @@
+# Config for testing nodes
+testing:
+    custom_nodes: tests/inference/testing_nodes
+

tests/inference/test_execution.py

@@ -0,0 +1,461 @@
+from io import BytesIO
+import numpy
+from PIL import Image
+import pytest
+from pytest import fixture
+import time
+import torch
+from typing import Union, Dict
+import json
+import subprocess
+import websocket #NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
+import uuid
+import urllib.request
+import urllib.parse
+import urllib.error
+from comfy_execution.graph_utils import GraphBuilder, Node
+
+class RunResult:
+    def __init__(self, prompt_id: str):
+        self.outputs: Dict[str,Dict] = {}
+        self.runs: Dict[str,bool] = {}
+        self.prompt_id: str = prompt_id
+
+    def get_output(self, node: Node):
+        return self.outputs.get(node.id, None)
+
+    def did_run(self, node: Node):
+        return self.runs.get(node.id, False)
+
+    def get_images(self, node: Node):
+        output = self.get_output(node)
+        if output is None:
+            return []
+        return output.get('image_objects', [])
+
+    def get_prompt_id(self):
+        return self.prompt_id
+
+class ComfyClient:
+    def __init__(self):
+        self.test_name = ""
+
+    def connect(self, 
+                    listen:str = '127.0.0.1', 
+                    port:Union[str,int] = 8188,
+                    client_id: str = str(uuid.uuid4())
+                    ):
+        self.client_id = client_id
+        self.server_address = f"{listen}:{port}"
+        ws = websocket.WebSocket()
+        ws.connect("ws://{}/ws?clientId={}".format(self.server_address, self.client_id))
+        self.ws = ws
+
+    def queue_prompt(self, prompt):
+        p = {"prompt": prompt, "client_id": self.client_id}
+        data = json.dumps(p).encode('utf-8')
+        req =  urllib.request.Request("http://{}/prompt".format(self.server_address), data=data)
+        return json.loads(urllib.request.urlopen(req).read())
+
+    def get_image(self, filename, subfolder, folder_type):
+        data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
+        url_values = urllib.parse.urlencode(data)
+        with urllib.request.urlopen("http://{}/view?{}".format(self.server_address, url_values)) as response:
+            return response.read()
+
+    def get_history(self, prompt_id):
+        with urllib.request.urlopen("http://{}/history/{}".format(self.server_address, prompt_id)) as response:
+            return json.loads(response.read())
+
+    def set_test_name(self, name):
+        self.test_name = name
+
+    def run(self, graph):
+        prompt = graph.finalize()
+        for node in graph.nodes.values():
+            if node.class_type == 'SaveImage':
+                node.inputs['filename_prefix'] = self.test_name
+
+        prompt_id = self.queue_prompt(prompt)['prompt_id']
+        result = RunResult(prompt_id)
+        while True:
+            out = self.ws.recv()
+            if isinstance(out, str):
+                message = json.loads(out)
+                if message['type'] == 'executing':
+                    data = message['data']
+                    if data['prompt_id'] != prompt_id:
+                        continue
+                    if data['node'] is None:
+                        break
+                    result.runs[data['node']] = True
+                elif message['type'] == 'execution_error':
+                    raise Exception(message['data'])
+                elif message['type'] == 'execution_cached':
+                    pass # Probably want to store this off for testing
+
+        history = self.get_history(prompt_id)[prompt_id]
+        for o in history['outputs']:
+            for node_id in history['outputs']:
+                node_output = history['outputs'][node_id]
+                result.outputs[node_id] = node_output
+                if 'images' in node_output:
+                    images_output = []
+                    for image in node_output['images']:
+                        image_data = self.get_image(image['filename'], image['subfolder'], image['type'])
+                        image_obj = Image.open(BytesIO(image_data))
+                        images_output.append(image_obj)
+                    node_output['image_objects'] = images_output
+
+        return result
+
+#
+# Loop through these variables
+#
+@pytest.mark.execution
+class TestExecution:
+    #
+    # Initialize server and client
+    #
+    @fixture(scope="class", autouse=True, params=[
+        # (use_lru, lru_size)
+        (False, 0),
+        (True, 0),
+        (True, 100),
+    ])
+    def _server(self, args_pytest, request):
+        # Start server
+        pargs = [
+            'python','main.py', 
+            '--output-directory', args_pytest["output_dir"],
+            '--listen', args_pytest["listen"],
+            '--port', str(args_pytest["port"]),
+            '--extra-model-paths-config', 'tests/inference/extra_model_paths.yaml',
+        ]
+        use_lru, lru_size = request.param
+        if use_lru:
+            pargs += ['--cache-lru', str(lru_size)]
+        print("Running server with args:", pargs)
+        p = subprocess.Popen(pargs)
+        yield
+        p.kill()
+        torch.cuda.empty_cache()
+
+    def start_client(self, listen:str, port:int):
+        # Start client
+        comfy_client = ComfyClient()
+        # Connect to server (with retries)
+        n_tries = 5
+        for i in range(n_tries):
+            time.sleep(4)
+            try:
+                comfy_client.connect(listen=listen, port=port)
+            except ConnectionRefusedError as e:
+                print(e)
+                print(f"({i+1}/{n_tries}) Retrying...")
+            else:
+                break
+        return comfy_client
+
+    @fixture(scope="class", autouse=True)
+    def shared_client(self, args_pytest, _server):
+        client = self.start_client(args_pytest["listen"], args_pytest["port"])
+        yield client
+        del client
+        torch.cuda.empty_cache()
+
+    @fixture
+    def client(self, shared_client, request):
+        shared_client.set_test_name(f"execution[{request.node.name}]")
+        yield shared_client
+
+    @fixture
+    def builder(self, request):
+        yield GraphBuilder(prefix=request.node.name)
+
+    def test_lazy_input(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
+        mask = g.node("StubMask", value=0.0, height=512, width=512, batch_size=1)
+
+        lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
+        output = g.node("SaveImage", images=lazy_mix.out(0))
+        result = client.run(g)
+
+        result_image = result.get_images(output)[0]
+        assert numpy.array(result_image).any() == 0, "Image should be black"
+        assert result.did_run(input1)
+        assert not result.did_run(input2)
+        assert result.did_run(mask)
+        assert result.did_run(lazy_mix)
+
+    def test_full_cache(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        input2 = g.node("StubImage", content="NOISE", height=512, width=512, batch_size=1)
+        mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
+
+        lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
+        g.node("SaveImage", images=lazy_mix.out(0))
+
+        client.run(g)
+        result2 = client.run(g)
+        for node_id, node in g.nodes.items():
+            assert not result2.did_run(node), f"Node {node_id} ran, but should have been cached"
+
+    def test_partial_cache(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        input2 = g.node("StubImage", content="NOISE", height=512, width=512, batch_size=1)
+        mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
+
+        lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
+        g.node("SaveImage", images=lazy_mix.out(0))
+
+        client.run(g)
+        mask.inputs['value'] = 0.4
+        result2 = client.run(g)
+        assert not result2.did_run(input1), "Input1 should have been cached"
+        assert not result2.did_run(input2), "Input2 should have been cached"
+
+    def test_error(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        # Different size of the two images
+        input2 = g.node("StubImage", content="NOISE", height=256, width=256, batch_size=1)
+        mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
+
+        lazy_mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
+        g.node("SaveImage", images=lazy_mix.out(0))
+
+        try:
+            client.run(g)
+            assert False, "Should have raised an error"
+        except Exception as e:
+            assert 'prompt_id' in e.args[0], f"Did not get back a proper error message: {e}"
+
+    @pytest.mark.parametrize("test_value, expect_error", [
+        (5, True),
+        ("foo", True),
+        (5.0, False),
+    ])
+    def test_validation_error_literal(self, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        validation1 = g.node("TestCustomValidation1", input1=test_value, input2=3.0)
+        g.node("SaveImage", images=validation1.out(0))
+
+        if expect_error:
+            with pytest.raises(urllib.error.HTTPError):
+                client.run(g)
+        else:
+            client.run(g)
+
+    @pytest.mark.parametrize("test_type, test_value", [
+        ("StubInt", 5),
+        ("StubFloat", 5.0)
+    ])
+    def test_validation_error_edge1(self, test_type, test_value, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        stub = g.node(test_type, value=test_value)
+        validation1 = g.node("TestCustomValidation1", input1=stub.out(0), input2=3.0)
+        g.node("SaveImage", images=validation1.out(0))
+
+        with pytest.raises(urllib.error.HTTPError):
+            client.run(g)
+
+    @pytest.mark.parametrize("test_type, test_value, expect_error", [
+        ("StubInt", 5, True),
+        ("StubFloat", 5.0, False)
+    ])
+    def test_validation_error_edge2(self, test_type, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        stub = g.node(test_type, value=test_value)
+        validation2 = g.node("TestCustomValidation2", input1=stub.out(0), input2=3.0)
+        g.node("SaveImage", images=validation2.out(0))
+
+        if expect_error:
+            with pytest.raises(urllib.error.HTTPError):
+                client.run(g)
+        else:
+            client.run(g)
+
+    @pytest.mark.parametrize("test_type, test_value, expect_error", [
+        ("StubInt", 5, True),
+        ("StubFloat", 5.0, False)
+    ])
+    def test_validation_error_edge3(self, test_type, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        stub = g.node(test_type, value=test_value)
+        validation3 = g.node("TestCustomValidation3", input1=stub.out(0), input2=3.0)
+        g.node("SaveImage", images=validation3.out(0))
+
+        if expect_error:
+            with pytest.raises(urllib.error.HTTPError):
+                client.run(g)
+        else:
+            client.run(g)
+
+    @pytest.mark.parametrize("test_type, test_value, expect_error", [
+        ("StubInt", 5, True),
+        ("StubFloat", 5.0, False)
+    ])
+    def test_validation_error_edge4(self, test_type, test_value, expect_error, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        stub = g.node(test_type, value=test_value)
+        validation4 = g.node("TestCustomValidation4", input1=stub.out(0), input2=3.0)
+        g.node("SaveImage", images=validation4.out(0))
+
+        if expect_error:
+            with pytest.raises(urllib.error.HTTPError):
+                client.run(g)
+        else:
+            client.run(g)
+
+    @pytest.mark.parametrize("test_value1, test_value2, expect_error", [
+        (0.0, 0.5, False),
+        (0.0, 5.0, False),
+        (0.0, 7.0, True)
+    ])
+    def test_validation_error_kwargs(self, test_value1, test_value2, expect_error, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        validation5 = g.node("TestCustomValidation5", input1=test_value1, input2=test_value2)
+        g.node("SaveImage", images=validation5.out(0))
+
+        if expect_error:
+            with pytest.raises(urllib.error.HTTPError):
+                client.run(g)
+        else:
+            client.run(g)
+
+    def test_cycle_error(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
+        mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
+
+        lazy_mix1 = g.node("TestLazyMixImages", image1=input1.out(0), mask=mask.out(0))
+        lazy_mix2 = g.node("TestLazyMixImages", image1=lazy_mix1.out(0), image2=input2.out(0), mask=mask.out(0))
+        g.node("SaveImage", images=lazy_mix2.out(0))
+
+        # When the cycle exists on initial submission, it should raise a validation error
+        with pytest.raises(urllib.error.HTTPError):
+            client.run(g)
+
+    def test_dynamic_cycle_error(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
+        generator = g.node("TestDynamicDependencyCycle", input1=input1.out(0), input2=input2.out(0))
+        g.node("SaveImage", images=generator.out(0))
+
+        # When the cycle is in a graph that is generated dynamically, it should raise a runtime error
+        try:
+            client.run(g)
+            assert False, "Should have raised an error"
+        except Exception as e:
+            assert 'prompt_id' in e.args[0], f"Did not get back a proper error message: {e}"
+            assert e.args[0]['node_id'] == generator.id, "Error should have been on the generator node"
+
+    def test_custom_is_changed(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        # Creating the nodes in this specific order previously caused a bug
+        save = g.node("SaveImage")
+        is_changed = g.node("TestCustomIsChanged", should_change=False)
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+
+        save.set_input('images', is_changed.out(0))
+        is_changed.set_input('image', input1.out(0))
+
+        result1 = client.run(g)
+        result2 = client.run(g)
+        is_changed.set_input('should_change', True)
+        result3 = client.run(g)
+        result4 = client.run(g)
+        assert result1.did_run(is_changed), "is_changed should have been run"
+        assert not result2.did_run(is_changed), "is_changed should have been cached"
+        assert result3.did_run(is_changed), "is_changed should have been re-run"
+        assert result4.did_run(is_changed), "is_changed should not have been cached"
+
+    def test_undeclared_inputs(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
+        input3 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        input4 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        average = g.node("TestVariadicAverage", input1=input1.out(0), input2=input2.out(0), input3=input3.out(0), input4=input4.out(0))
+        output = g.node("SaveImage", images=average.out(0))
+
+        result = client.run(g)
+        result_image = result.get_images(output)[0]
+        expected = 255 // 4
+        assert numpy.array(result_image).min() == expected and numpy.array(result_image).max() == expected, "Image should be grey"
+
+    def test_for_loop(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        iterations = 4
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
+        is_changed = g.node("TestCustomIsChanged", should_change=True, image=input2.out(0))
+        for_open = g.node("TestForLoopOpen", remaining=iterations, initial_value1=is_changed.out(0))
+        average = g.node("TestVariadicAverage", input1=input1.out(0), input2=for_open.out(2))
+        for_close = g.node("TestForLoopClose", flow_control=for_open.out(0), initial_value1=average.out(0))
+        output = g.node("SaveImage", images=for_close.out(0))
+
+        for iterations in range(1, 5):
+            for_open.set_input('remaining', iterations)
+            result = client.run(g)
+            result_image = result.get_images(output)[0]
+            expected = 255 // (2 ** iterations)
+            assert numpy.array(result_image).min() == expected and numpy.array(result_image).max() == expected, "Image should be grey"
+            assert result.did_run(is_changed)
+
+    def test_mixed_expansion_returns(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        val_list = g.node("TestMakeListNode", value1=0.1, value2=0.2, value3=0.3)
+        mixed = g.node("TestMixedExpansionReturns", input1=val_list.out(0))
+        output_dynamic = g.node("SaveImage", images=mixed.out(0))
+        output_literal = g.node("SaveImage", images=mixed.out(1))
+
+        result = client.run(g)
+        images_dynamic = result.get_images(output_dynamic)
+        assert len(images_dynamic) == 3, "Should have 2 images"
+        assert numpy.array(images_dynamic[0]).min() == 25 and numpy.array(images_dynamic[0]).max() == 25, "First image should be 0.1"
+        assert numpy.array(images_dynamic[1]).min() == 51 and numpy.array(images_dynamic[1]).max() == 51, "Second image should be 0.2"
+        assert numpy.array(images_dynamic[2]).min() == 76 and numpy.array(images_dynamic[2]).max() == 76, "Third image should be 0.3"
+
+        images_literal = result.get_images(output_literal)
+        assert len(images_literal) == 3, "Should have 2 images"
+        for i in range(3):
+            assert numpy.array(images_literal[i]).min() == 255 and numpy.array(images_literal[i]).max() == 255, "All images should be white"
+
+    def test_mixed_lazy_results(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        val_list = g.node("TestMakeListNode", value1=0.0, value2=0.5, value3=1.0)
+        mask = g.node("StubMask", value=val_list.out(0), height=512, width=512, batch_size=1)
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+        input2 = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
+        mix = g.node("TestLazyMixImages", image1=input1.out(0), image2=input2.out(0), mask=mask.out(0))
+        rebatch = g.node("RebatchImages", images=mix.out(0), batch_size=3)
+        output = g.node("SaveImage", images=rebatch.out(0))
+
+        result = client.run(g)
+        images = result.get_images(output)
+        assert len(images) == 3, "Should have 3 image"
+        assert numpy.array(images[0]).min() == 0 and numpy.array(images[0]).max() == 0, "First image should be 0.0"
+        assert numpy.array(images[1]).min() == 127 and numpy.array(images[1]).max() == 127, "Second image should be 0.5"
+        assert numpy.array(images[2]).min() == 255 and numpy.array(images[2]).max() == 255, "Third image should be 1.0"
+
+    def test_output_reuse(self, client: ComfyClient, builder: GraphBuilder):
+        g = builder
+        input1 = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+
+        output1 = g.node("PreviewImage", images=input1.out(0))
+        output2 = g.node("PreviewImage", images=input1.out(0))
+
+        result = client.run(g)
+        images1 = result.get_images(output1)
+        images2 = result.get_images(output2)
+        assert len(images1) == 1, "Should have 1 image"
+        assert len(images2) == 1, "Should have 1 image"
+

tests/inference/testing_nodes/testing-pack/__init__.py

@@ -0,0 +1,23 @@
+from .specific_tests import TEST_NODE_CLASS_MAPPINGS, TEST_NODE_DISPLAY_NAME_MAPPINGS
+from .flow_control import FLOW_CONTROL_NODE_CLASS_MAPPINGS, FLOW_CONTROL_NODE_DISPLAY_NAME_MAPPINGS
+from .util import UTILITY_NODE_CLASS_MAPPINGS, UTILITY_NODE_DISPLAY_NAME_MAPPINGS
+from .conditions import CONDITION_NODE_CLASS_MAPPINGS, CONDITION_NODE_DISPLAY_NAME_MAPPINGS
+from .stubs import TEST_STUB_NODE_CLASS_MAPPINGS, TEST_STUB_NODE_DISPLAY_NAME_MAPPINGS
+
+# NODE_CLASS_MAPPINGS = GENERAL_NODE_CLASS_MAPPINGS.update(COMPONENT_NODE_CLASS_MAPPINGS)
+# NODE_DISPLAY_NAME_MAPPINGS = GENERAL_NODE_DISPLAY_NAME_MAPPINGS.update(COMPONENT_NODE_DISPLAY_NAME_MAPPINGS)
+
+NODE_CLASS_MAPPINGS = {}
+NODE_CLASS_MAPPINGS.update(TEST_NODE_CLASS_MAPPINGS)
+NODE_CLASS_MAPPINGS.update(FLOW_CONTROL_NODE_CLASS_MAPPINGS)
+NODE_CLASS_MAPPINGS.update(UTILITY_NODE_CLASS_MAPPINGS)
+NODE_CLASS_MAPPINGS.update(CONDITION_NODE_CLASS_MAPPINGS)
+NODE_CLASS_MAPPINGS.update(TEST_STUB_NODE_CLASS_MAPPINGS)
+
+NODE_DISPLAY_NAME_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS.update(TEST_NODE_DISPLAY_NAME_MAPPINGS)
+NODE_DISPLAY_NAME_MAPPINGS.update(FLOW_CONTROL_NODE_DISPLAY_NAME_MAPPINGS)
+NODE_DISPLAY_NAME_MAPPINGS.update(UTILITY_NODE_DISPLAY_NAME_MAPPINGS)
+NODE_DISPLAY_NAME_MAPPINGS.update(CONDITION_NODE_DISPLAY_NAME_MAPPINGS)
+NODE_DISPLAY_NAME_MAPPINGS.update(TEST_STUB_NODE_DISPLAY_NAME_MAPPINGS)
+

tests/inference/testing_nodes/testing-pack/conditions.py

@@ -0,0 +1,194 @@
+import re
+import torch
+
+class TestIntConditions:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "a": ("INT", {"default": 0, "min": -0xffffffffffffffff, "max": 0xffffffffffffffff, "step": 1}),
+                "b": ("INT", {"default": 0, "min": -0xffffffffffffffff, "max": 0xffffffffffffffff, "step": 1}),
+                "operation": (["==", "!=", "<", ">", "<=", ">="],),
+            },
+        }
+
+    RETURN_TYPES = ("BOOLEAN",)
+    FUNCTION = "int_condition"
+
+    CATEGORY = "Testing/Logic"
+
+    def int_condition(self, a, b, operation):
+        if operation == "==":
+            return (a == b,)
+        elif operation == "!=":
+            return (a != b,)
+        elif operation == "<":
+            return (a < b,)
+        elif operation == ">":
+            return (a > b,)
+        elif operation == "<=":
+            return (a <= b,)
+        elif operation == ">=":
+            return (a >= b,)
+
+
+class TestFloatConditions:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "a": ("FLOAT", {"default": 0, "min": -999999999999.0, "max": 999999999999.0, "step": 1}),
+                "b": ("FLOAT", {"default": 0, "min": -999999999999.0, "max": 999999999999.0, "step": 1}),
+                "operation": (["==", "!=", "<", ">", "<=", ">="],),
+            },
+        }
+
+    RETURN_TYPES = ("BOOLEAN",)
+    FUNCTION = "float_condition"
+
+    CATEGORY = "Testing/Logic"
+
+    def float_condition(self, a, b, operation):
+        if operation == "==":
+            return (a == b,)
+        elif operation == "!=":
+            return (a != b,)
+        elif operation == "<":
+            return (a < b,)
+        elif operation == ">":
+            return (a > b,)
+        elif operation == "<=":
+            return (a <= b,)
+        elif operation == ">=":
+            return (a >= b,)
+
+class TestStringConditions:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "a": ("STRING", {"multiline": False}),
+                "b": ("STRING", {"multiline": False}),
+                "operation": (["a == b", "a != b", "a IN b", "a MATCH REGEX(b)", "a BEGINSWITH b", "a ENDSWITH b"],),
+                "case_sensitive": ("BOOLEAN", {"default": True}),
+            },
+        }
+
+    RETURN_TYPES = ("BOOLEAN",)
+    FUNCTION = "string_condition"
+
+    CATEGORY = "Testing/Logic"
+
+    def string_condition(self, a, b, operation, case_sensitive):
+        if not case_sensitive:
+            a = a.lower()
+            b = b.lower()
+
+        if operation == "a == b":
+            return (a == b,)
+        elif operation == "a != b":
+            return (a != b,)
+        elif operation == "a IN b":
+            return (a in b,)
+        elif operation == "a MATCH REGEX(b)":
+            try:
+                return (re.match(b, a) is not None,)
+            except:
+                return (False,)
+        elif operation == "a BEGINSWITH b":
+            return (a.startswith(b),)
+        elif operation == "a ENDSWITH b":
+            return (a.endswith(b),)
+
+class TestToBoolNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "value": ("*",),
+            },
+            "optional": {
+                "invert": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = ("BOOLEAN",)
+    FUNCTION = "to_bool"
+
+    CATEGORY = "Testing/Logic"
+
+    def to_bool(self, value, invert = False):
+        if isinstance(value, torch.Tensor):
+            if value.max().item() == 0 and value.min().item() == 0:
+                result = False
+            else:
+                result = True
+        else:
+            try:
+                result = bool(value)
+            except:
+                # Can't convert it? Well then it's something or other. I dunno, I'm not a Python programmer.
+                result = True
+
+        if invert:
+            result = not result
+
+        return (result,)
+
+class TestBoolOperationNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "a": ("BOOLEAN",),
+                "b": ("BOOLEAN",),
+                "op": (["a AND b", "a OR b", "a XOR b", "NOT a"],),
+            },
+        }
+
+    RETURN_TYPES = ("BOOLEAN",)
+    FUNCTION = "bool_operation"
+
+    CATEGORY = "Testing/Logic"
+
+    def bool_operation(self, a, b, op):
+        if op == "a AND b":
+            return (a and b,)
+        elif op == "a OR b":
+            return (a or b,)
+        elif op == "a XOR b":
+            return (a ^ b,)
+        elif op == "NOT a":
+            return (not a,)
+
+
+CONDITION_NODE_CLASS_MAPPINGS = {
+    "TestIntConditions": TestIntConditions,
+    "TestFloatConditions": TestFloatConditions,
+    "TestStringConditions": TestStringConditions,
+    "TestToBoolNode": TestToBoolNode,
+    "TestBoolOperationNode": TestBoolOperationNode,
+}
+
+CONDITION_NODE_DISPLAY_NAME_MAPPINGS = {
+    "TestIntConditions": "Int Condition",
+    "TestFloatConditions": "Float Condition",
+    "TestStringConditions": "String Condition",
+    "TestToBoolNode": "To Bool",
+    "TestBoolOperationNode": "Bool Operation",
+}

tests/inference/testing_nodes/testing-pack/flow_control.py

@@ -0,0 +1,173 @@
+from comfy_execution.graph_utils import GraphBuilder, is_link
+from comfy_execution.graph import ExecutionBlocker
+from .tools import VariantSupport
+
+NUM_FLOW_SOCKETS = 5
+@VariantSupport()
+class TestWhileLoopOpen:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        inputs = {
+            "required": {
+                "condition": ("BOOLEAN", {"default": True}),
+            },
+            "optional": {
+            },
+        }
+        for i in range(NUM_FLOW_SOCKETS):
+            inputs["optional"][f"initial_value{i}"] = ("*",)
+        return inputs
+
+    RETURN_TYPES = tuple(["FLOW_CONTROL"] + ["*"] * NUM_FLOW_SOCKETS)
+    RETURN_NAMES = tuple(["FLOW_CONTROL"] + [f"value{i}" for i in range(NUM_FLOW_SOCKETS)])
+    FUNCTION = "while_loop_open"
+
+    CATEGORY = "Testing/Flow"
+
+    def while_loop_open(self, condition, **kwargs):
+        values = []
+        for i in range(NUM_FLOW_SOCKETS):
+            values.append(kwargs.get(f"initial_value{i}", None))
+        return tuple(["stub"] + values)
+
+@VariantSupport()
+class TestWhileLoopClose:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        inputs = {
+            "required": {
+                "flow_control": ("FLOW_CONTROL", {"rawLink": True}),
+                "condition": ("BOOLEAN", {"forceInput": True}),
+            },
+            "optional": {
+            },
+            "hidden": {
+                "dynprompt": "DYNPROMPT",
+                "unique_id": "UNIQUE_ID",
+            }
+        }
+        for i in range(NUM_FLOW_SOCKETS):
+            inputs["optional"][f"initial_value{i}"] = ("*",)
+        return inputs
+
+    RETURN_TYPES = tuple(["*"] * NUM_FLOW_SOCKETS)
+    RETURN_NAMES = tuple([f"value{i}" for i in range(NUM_FLOW_SOCKETS)])
+    FUNCTION = "while_loop_close"
+
+    CATEGORY = "Testing/Flow"
+
+    def explore_dependencies(self, node_id, dynprompt, upstream):
+        node_info = dynprompt.get_node(node_id)
+        if "inputs" not in node_info:
+            return
+        for k, v in node_info["inputs"].items():
+            if is_link(v):
+                parent_id = v[0]
+                if parent_id not in upstream:
+                    upstream[parent_id] = []
+                    self.explore_dependencies(parent_id, dynprompt, upstream)
+                upstream[parent_id].append(node_id)
+
+    def collect_contained(self, node_id, upstream, contained):
+        if node_id not in upstream:
+            return
+        for child_id in upstream[node_id]:
+            if child_id not in contained:
+                contained[child_id] = True
+                self.collect_contained(child_id, upstream, contained)
+
+
+    def while_loop_close(self, flow_control, condition, dynprompt=None, unique_id=None, **kwargs):
+        assert dynprompt is not None
+        if not condition:
+            # We're done with the loop
+            values = []
+            for i in range(NUM_FLOW_SOCKETS):
+                values.append(kwargs.get(f"initial_value{i}", None))
+            return tuple(values)
+
+        # We want to loop
+        upstream = {}
+        # Get the list of all nodes between the open and close nodes
+        self.explore_dependencies(unique_id, dynprompt, upstream)
+
+        contained = {}
+        open_node = flow_control[0]
+        self.collect_contained(open_node, upstream, contained)
+        contained[unique_id] = True
+        contained[open_node] = True
+
+        # We'll use the default prefix, but to avoid having node names grow exponentially in size,
+        # we'll use "Recurse" for the name of the recursively-generated copy of this node.
+        graph = GraphBuilder()
+        for node_id in contained:
+            original_node = dynprompt.get_node(node_id)
+            node = graph.node(original_node["class_type"], "Recurse" if node_id == unique_id else node_id)
+            node.set_override_display_id(node_id)
+        for node_id in contained:
+            original_node = dynprompt.get_node(node_id)
+            node = graph.lookup_node("Recurse" if node_id == unique_id else node_id)
+            assert node is not None
+            for k, v in original_node["inputs"].items():
+                if is_link(v) and v[0] in contained:
+                    parent = graph.lookup_node(v[0])
+                    assert parent is not None
+                    node.set_input(k, parent.out(v[1]))
+                else:
+                    node.set_input(k, v)
+        new_open = graph.lookup_node(open_node)
+        assert new_open is not None
+        for i in range(NUM_FLOW_SOCKETS):
+            key = f"initial_value{i}"
+            new_open.set_input(key, kwargs.get(key, None))
+        my_clone = graph.lookup_node("Recurse")
+        assert my_clone is not None
+        result = map(lambda x: my_clone.out(x), range(NUM_FLOW_SOCKETS))
+        return {
+            "result": tuple(result),
+            "expand": graph.finalize(),
+        }
+
+@VariantSupport()
+class TestExecutionBlockerNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        inputs = {
+            "required": {
+                "input": ("*",),
+                "block": ("BOOLEAN",),
+                "verbose": ("BOOLEAN", {"default": False}),
+            },
+        }
+        return inputs
+
+    RETURN_TYPES = ("*",)
+    RETURN_NAMES = ("output",)
+    FUNCTION = "execution_blocker"
+
+    CATEGORY = "Testing/Flow"
+
+    def execution_blocker(self, input, block, verbose):
+        if block:
+            return (ExecutionBlocker("Blocked Execution" if verbose else None),)
+        return (input,)
+
+FLOW_CONTROL_NODE_CLASS_MAPPINGS = {
+    "TestWhileLoopOpen": TestWhileLoopOpen,
+    "TestWhileLoopClose": TestWhileLoopClose,
+    "TestExecutionBlocker": TestExecutionBlockerNode,
+}
+FLOW_CONTROL_NODE_DISPLAY_NAME_MAPPINGS = {
+    "TestWhileLoopOpen": "While Loop Open",
+    "TestWhileLoopClose": "While Loop Close",
+    "TestExecutionBlocker": "Execution Blocker",
+}

tests/inference/testing_nodes/testing-pack/specific_tests.py

@@ -0,0 +1,335 @@
+import torch
+from .tools import VariantSupport
+from comfy_execution.graph_utils import GraphBuilder
+
+class TestLazyMixImages:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image1": ("IMAGE",{"lazy": True}),
+                "image2": ("IMAGE",{"lazy": True}),
+                "mask": ("MASK",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "mix"
+
+    CATEGORY = "Testing/Nodes"
+
+    def check_lazy_status(self, mask, image1, image2):
+        mask_min = mask.min()
+        mask_max = mask.max()
+        needed = []
+        if image1 is None and (mask_min != 1.0 or mask_max != 1.0):
+            needed.append("image1")
+        if image2 is None and (mask_min != 0.0 or mask_max != 0.0):
+            needed.append("image2")
+        return needed
+
+    # Not trying to handle different batch sizes here just to keep the demo simple
+    def mix(self, mask, image1, image2):
+        mask_min = mask.min()
+        mask_max = mask.max()
+        if mask_min == 0.0 and mask_max == 0.0:
+            return (image1,)
+        elif mask_min == 1.0 and mask_max == 1.0:
+            return (image2,)
+
+        if len(mask.shape) == 2:
+            mask = mask.unsqueeze(0)
+        if len(mask.shape) == 3:
+            mask = mask.unsqueeze(3)
+        if mask.shape[3] < image1.shape[3]:
+            mask = mask.repeat(1, 1, 1, image1.shape[3])
+
+        result = image1 * (1. - mask) + image2 * mask,
+        return (result[0],)
+
+class TestVariadicAverage:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input1": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "variadic_average"
+
+    CATEGORY = "Testing/Nodes"
+
+    def variadic_average(self, input1, **kwargs):
+        inputs = [input1]
+        while 'input' + str(len(inputs) + 1) in kwargs:
+            inputs.append(kwargs['input' + str(len(inputs) + 1)])
+        return (torch.stack(inputs).mean(dim=0),)
+
+
+class TestCustomIsChanged:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+            },
+            "optional": {
+                "should_change": ("BOOL", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "custom_is_changed"
+
+    CATEGORY = "Testing/Nodes"
+
+    def custom_is_changed(self, image, should_change=False):
+        return (image,)
+    
+    @classmethod
+    def IS_CHANGED(cls, should_change=False, *args, **kwargs):
+        if should_change:
+            return float("NaN")
+        else:
+            return False
+
+class TestCustomValidation1:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input1": ("IMAGE,FLOAT",),
+                "input2": ("IMAGE,FLOAT",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "custom_validation1"
+
+    CATEGORY = "Testing/Nodes"
+
+    def custom_validation1(self, input1, input2):
+        if isinstance(input1, float) and isinstance(input2, float):
+            result = torch.ones([1, 512, 512, 3]) * input1 * input2
+        else:
+            result = input1 * input2
+        return (result,)
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, input1=None, input2=None):
+        if input1 is not None:
+            if not isinstance(input1, (torch.Tensor, float)):
+                return f"Invalid type of input1: {type(input1)}"
+        if input2 is not None:
+            if not isinstance(input2, (torch.Tensor, float)):
+                return f"Invalid type of input2: {type(input2)}"
+
+        return True
+
+class TestCustomValidation2:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input1": ("IMAGE,FLOAT",),
+                "input2": ("IMAGE,FLOAT",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "custom_validation2"
+
+    CATEGORY = "Testing/Nodes"
+
+    def custom_validation2(self, input1, input2):
+        if isinstance(input1, float) and isinstance(input2, float):
+            result = torch.ones([1, 512, 512, 3]) * input1 * input2
+        else:
+            result = input1 * input2
+        return (result,)
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, input_types, input1=None, input2=None):
+        if input1 is not None:
+            if not isinstance(input1, (torch.Tensor, float)):
+                return f"Invalid type of input1: {type(input1)}"
+        if input2 is not None:
+            if not isinstance(input2, (torch.Tensor, float)):
+                return f"Invalid type of input2: {type(input2)}"
+
+        if 'input1' in input_types:
+            if input_types['input1'] not in ["IMAGE", "FLOAT"]:
+                return f"Invalid type of input1: {input_types['input1']}"
+        if 'input2' in input_types:
+            if input_types['input2'] not in ["IMAGE", "FLOAT"]:
+                return f"Invalid type of input2: {input_types['input2']}"
+
+        return True
+
+@VariantSupport()
+class TestCustomValidation3:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input1": ("IMAGE,FLOAT",),
+                "input2": ("IMAGE,FLOAT",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "custom_validation3"
+
+    CATEGORY = "Testing/Nodes"
+
+    def custom_validation3(self, input1, input2):
+        if isinstance(input1, float) and isinstance(input2, float):
+            result = torch.ones([1, 512, 512, 3]) * input1 * input2
+        else:
+            result = input1 * input2
+        return (result,)
+
+class TestCustomValidation4:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input1": ("FLOAT",),
+                "input2": ("FLOAT",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "custom_validation4"
+
+    CATEGORY = "Testing/Nodes"
+
+    def custom_validation4(self, input1, input2):
+        result = torch.ones([1, 512, 512, 3]) * input1 * input2
+        return (result,)
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, input1, input2):
+        if input1 is not None:
+            if not isinstance(input1, float):
+                return f"Invalid type of input1: {type(input1)}"
+        if input2 is not None:
+            if not isinstance(input2, float):
+                return f"Invalid type of input2: {type(input2)}"
+
+        return True
+
+class TestCustomValidation5:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input1": ("FLOAT", {"min": 0.0, "max": 1.0}),
+                "input2": ("FLOAT", {"min": 0.0, "max": 1.0}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "custom_validation5"
+
+    CATEGORY = "Testing/Nodes"
+
+    def custom_validation5(self, input1, input2):
+        value = input1 * input2
+        return (torch.ones([1, 512, 512, 3]) * value,)
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, **kwargs):
+        if kwargs['input2'] == 7.0:
+            return "7s are not allowed. I've never liked 7s."
+        return True
+
+class TestDynamicDependencyCycle:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input1": ("IMAGE",),
+                "input2": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "dynamic_dependency_cycle"
+
+    CATEGORY = "Testing/Nodes"
+
+    def dynamic_dependency_cycle(self, input1, input2):
+        g = GraphBuilder()
+        mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
+        mix1 = g.node("TestLazyMixImages", image1=input1, mask=mask.out(0))
+        mix2 = g.node("TestLazyMixImages", image1=mix1.out(0), image2=input2, mask=mask.out(0))
+
+        # Create the cyle
+        mix1.set_input("image2", mix2.out(0))
+
+        return {
+            "result": (mix2.out(0),),
+            "expand": g.finalize(),
+        }
+
+class TestMixedExpansionReturns:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input1": ("FLOAT",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE","IMAGE")
+    FUNCTION = "mixed_expansion_returns"
+
+    CATEGORY = "Testing/Nodes"
+
+    def mixed_expansion_returns(self, input1):
+        white_image = torch.ones([1, 512, 512, 3])
+        if input1 <= 0.1:
+            return (torch.ones([1, 512, 512, 3]) * 0.1, white_image)
+        elif input1 <= 0.2:
+            return {
+                "result": (torch.ones([1, 512, 512, 3]) * 0.2, white_image),
+            }
+        else:
+            g = GraphBuilder()
+            mask = g.node("StubMask", value=0.3, height=512, width=512, batch_size=1)
+            black = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
+            white = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
+            mix = g.node("TestLazyMixImages", image1=black.out(0), image2=white.out(0), mask=mask.out(0))
+            return {
+                "result": (mix.out(0), white_image),
+                "expand": g.finalize(),
+            }
+
+TEST_NODE_CLASS_MAPPINGS = {
+    "TestLazyMixImages": TestLazyMixImages,
+    "TestVariadicAverage": TestVariadicAverage,
+    "TestCustomIsChanged": TestCustomIsChanged,
+    "TestCustomValidation1": TestCustomValidation1,
+    "TestCustomValidation2": TestCustomValidation2,
+    "TestCustomValidation3": TestCustomValidation3,
+    "TestCustomValidation4": TestCustomValidation4,
+    "TestCustomValidation5": TestCustomValidation5,
+    "TestDynamicDependencyCycle": TestDynamicDependencyCycle,
+    "TestMixedExpansionReturns": TestMixedExpansionReturns,
+}
+
+TEST_NODE_DISPLAY_NAME_MAPPINGS = {
+    "TestLazyMixImages": "Lazy Mix Images",
+    "TestVariadicAverage": "Variadic Average",
+    "TestCustomIsChanged": "Custom IsChanged",
+    "TestCustomValidation1": "Custom Validation 1",
+    "TestCustomValidation2": "Custom Validation 2",
+    "TestCustomValidation3": "Custom Validation 3",
+    "TestCustomValidation4": "Custom Validation 4",
+    "TestCustomValidation5": "Custom Validation 5",
+    "TestDynamicDependencyCycle": "Dynamic Dependency Cycle",
+    "TestMixedExpansionReturns": "Mixed Expansion Returns",
+}

tests/inference/testing_nodes/testing-pack/stubs.py

@@ -0,0 +1,105 @@
+import torch
+
+class StubImage:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "content": (['WHITE', 'BLACK', 'NOISE'],),
+                "height": ("INT", {"default": 512, "min": 1, "max": 1024 ** 3, "step": 1}),
+                "width": ("INT", {"default": 512, "min": 1, "max": 4096 ** 3, "step": 1}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 1024 ** 3, "step": 1}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "stub_image"
+
+    CATEGORY = "Testing/Stub Nodes"
+
+    def stub_image(self, content, height, width, batch_size):
+        if content == "WHITE":
+            return (torch.ones(batch_size, height, width, 3),)
+        elif content == "BLACK":
+            return (torch.zeros(batch_size, height, width, 3),)
+        elif content == "NOISE":
+            return (torch.rand(batch_size, height, width, 3),)
+
+class StubMask:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "value": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "height": ("INT", {"default": 512, "min": 1, "max": 1024 ** 3, "step": 1}),
+                "width": ("INT", {"default": 512, "min": 1, "max": 4096 ** 3, "step": 1}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 1024 ** 3, "step": 1}),
+            },
+        }
+
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "stub_mask"
+
+    CATEGORY = "Testing/Stub Nodes"
+
+    def stub_mask(self, value, height, width, batch_size):
+        return (torch.ones(batch_size, height, width) * value,)
+
+class StubInt:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "value": ("INT", {"default": 0, "min": -0xffffffff, "max": 0xffffffff, "step": 1}),
+            },
+        }
+
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "stub_int"
+
+    CATEGORY = "Testing/Stub Nodes"
+
+    def stub_int(self, value):
+        return (value,)
+
+class StubFloat:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "value": ("FLOAT", {"default": 0.0, "min": -1.0e38, "max": 1.0e38, "step": 0.01}),
+            },
+        }
+
+    RETURN_TYPES = ("FLOAT",)
+    FUNCTION = "stub_float"
+
+    CATEGORY = "Testing/Stub Nodes"
+
+    def stub_float(self, value):
+        return (value,)
+
+TEST_STUB_NODE_CLASS_MAPPINGS = {
+    "StubImage": StubImage,
+    "StubMask": StubMask,
+    "StubInt": StubInt,
+    "StubFloat": StubFloat,
+}
+TEST_STUB_NODE_DISPLAY_NAME_MAPPINGS = {
+    "StubImage": "Stub Image",
+    "StubMask": "Stub Mask",
+    "StubInt": "Stub Int",
+    "StubFloat": "Stub Float",
+}

tests/inference/testing_nodes/testing-pack/tools.py

@@ -0,0 +1,53 @@
+
+def MakeSmartType(t):
+    if isinstance(t, str):
+        return SmartType(t)
+    return t
+
+class SmartType(str):
+    def __ne__(self, other):
+        if self == "*" or other == "*":
+            return False
+        selfset = set(self.split(','))
+        otherset = set(other.split(','))
+        return not selfset.issubset(otherset)
+
+def VariantSupport():
+    def decorator(cls):
+        if hasattr(cls, "INPUT_TYPES"):
+            old_input_types = getattr(cls, "INPUT_TYPES")
+            def new_input_types(*args, **kwargs):
+                types = old_input_types(*args, **kwargs)
+                for category in ["required", "optional"]:
+                    if category not in types:
+                        continue
+                    for key, value in types[category].items():
+                        if isinstance(value, tuple):
+                            types[category][key] = (MakeSmartType(value[0]),) + value[1:]
+                return types
+            setattr(cls, "INPUT_TYPES", new_input_types)
+        if hasattr(cls, "RETURN_TYPES"):
+            old_return_types = cls.RETURN_TYPES
+            setattr(cls, "RETURN_TYPES", tuple(MakeSmartType(x) for x in old_return_types))
+        if hasattr(cls, "VALIDATE_INPUTS"):
+            # Reflection is used to determine what the function signature is, so we can't just change the function signature
+            raise NotImplementedError("VariantSupport does not support VALIDATE_INPUTS yet")
+        else:
+            def validate_inputs(input_types):
+                inputs = cls.INPUT_TYPES()
+                for key, value in input_types.items():
+                    if isinstance(value, SmartType):
+                        continue
+                    if "required" in inputs and key in inputs["required"]:
+                        expected_type = inputs["required"][key][0]
+                    elif "optional" in inputs and key in inputs["optional"]:
+                        expected_type = inputs["optional"][key][0]
+                    else:
+                        expected_type = None
+                    if expected_type is not None and MakeSmartType(value) != expected_type:
+                        return f"Invalid type of {key}: {value} (expected {expected_type})"
+                return True
+            setattr(cls, "VALIDATE_INPUTS", validate_inputs)
+        return cls
+    return decorator
+

tests/inference/testing_nodes/testing-pack/util.py

@@ -0,0 +1,364 @@
+from comfy_execution.graph_utils import GraphBuilder
+from .tools import VariantSupport
+
+@VariantSupport()
+class TestAccumulateNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "to_add": ("*",),
+            },
+            "optional": {
+                "accumulation": ("ACCUMULATION",),
+            },
+        }
+
+    RETURN_TYPES = ("ACCUMULATION",)
+    FUNCTION = "accumulate"
+
+    CATEGORY = "Testing/Lists"
+
+    def accumulate(self, to_add, accumulation = None):
+        if accumulation is None:
+            value = [to_add]
+        else:
+            value = accumulation["accum"] + [to_add]
+        return ({"accum": value},)
+
+@VariantSupport()
+class TestAccumulationHeadNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "accumulation": ("ACCUMULATION",),
+            },
+        }
+
+    RETURN_TYPES = ("ACCUMULATION", "*",)
+    FUNCTION = "accumulation_head"
+
+    CATEGORY = "Testing/Lists"
+
+    def accumulation_head(self, accumulation):
+        accum = accumulation["accum"]
+        if len(accum) == 0:
+            return (accumulation, None)
+        else:
+            return ({"accum": accum[1:]}, accum[0])
+
+class TestAccumulationTailNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "accumulation": ("ACCUMULATION",),
+            },
+        }
+
+    RETURN_TYPES = ("ACCUMULATION", "*",)
+    FUNCTION = "accumulation_tail"
+
+    CATEGORY = "Testing/Lists"
+
+    def accumulation_tail(self, accumulation):
+        accum = accumulation["accum"]
+        if len(accum) == 0:
+            return (None, accumulation)
+        else:
+            return ({"accum": accum[:-1]}, accum[-1])
+
+@VariantSupport()
+class TestAccumulationToListNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "accumulation": ("ACCUMULATION",),
+            },
+        }
+
+    RETURN_TYPES = ("*",)
+    OUTPUT_IS_LIST = (True,)
+
+    FUNCTION = "accumulation_to_list"
+
+    CATEGORY = "Testing/Lists"
+
+    def accumulation_to_list(self, accumulation):
+        return (accumulation["accum"],)
+
+@VariantSupport()
+class TestListToAccumulationNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "list": ("*",),
+            },
+        }
+
+    RETURN_TYPES = ("ACCUMULATION",)
+    INPUT_IS_LIST = (True,)
+
+    FUNCTION = "list_to_accumulation"
+
+    CATEGORY = "Testing/Lists"
+
+    def list_to_accumulation(self, list):
+        return ({"accum": list},)
+
+@VariantSupport()
+class TestAccumulationGetLengthNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "accumulation": ("ACCUMULATION",),
+            },
+        }
+
+    RETURN_TYPES = ("INT",)
+
+    FUNCTION = "accumlength"
+
+    CATEGORY = "Testing/Lists"
+
+    def accumlength(self, accumulation):
+        return (len(accumulation['accum']),)
+        
+@VariantSupport()
+class TestAccumulationGetItemNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "accumulation": ("ACCUMULATION",),
+                "index": ("INT", {"default":0, "step":1})
+            },
+        }
+
+    RETURN_TYPES = ("*",)
+
+    FUNCTION = "get_item"
+
+    CATEGORY = "Testing/Lists"
+
+    def get_item(self, accumulation, index):
+        return (accumulation['accum'][index],)
+        
+@VariantSupport()
+class TestAccumulationSetItemNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "accumulation": ("ACCUMULATION",),
+                "index": ("INT", {"default":0, "step":1}),
+                "value": ("*",),
+            },
+        }
+
+    RETURN_TYPES = ("ACCUMULATION",)
+
+    FUNCTION = "set_item"
+
+    CATEGORY = "Testing/Lists"
+
+    def set_item(self, accumulation, index, value):
+        new_accum = accumulation['accum'][:]
+        new_accum[index] = value
+        return ({"accum": new_accum},)
+
+class TestIntMathOperation:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "a": ("INT", {"default": 0, "min": -0xffffffffffffffff, "max": 0xffffffffffffffff, "step": 1}),
+                "b": ("INT", {"default": 0, "min": -0xffffffffffffffff, "max": 0xffffffffffffffff, "step": 1}),
+                "operation": (["add", "subtract", "multiply", "divide", "modulo", "power"],),
+            },
+        }
+
+    RETURN_TYPES = ("INT",)
+    FUNCTION = "int_math_operation"
+
+    CATEGORY = "Testing/Logic"
+
+    def int_math_operation(self, a, b, operation):
+        if operation == "add":
+            return (a + b,)
+        elif operation == "subtract":
+            return (a - b,)
+        elif operation == "multiply":
+            return (a * b,)
+        elif operation == "divide":
+            return (a // b,)
+        elif operation == "modulo":
+            return (a % b,)
+        elif operation == "power":
+            return (a ** b,)
+
+
+from .flow_control import NUM_FLOW_SOCKETS
+@VariantSupport()
+class TestForLoopOpen:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "remaining": ("INT", {"default": 1, "min": 0, "max": 100000, "step": 1}),
+            },
+            "optional": {
+                f"initial_value{i}": ("*",) for i in range(1, NUM_FLOW_SOCKETS)
+            },
+            "hidden": {
+                "initial_value0": ("*",)
+            }
+        }
+
+    RETURN_TYPES = tuple(["FLOW_CONTROL", "INT",] + ["*"] * (NUM_FLOW_SOCKETS-1))
+    RETURN_NAMES = tuple(["flow_control", "remaining"] + [f"value{i}" for i in range(1, NUM_FLOW_SOCKETS)])
+    FUNCTION = "for_loop_open"
+
+    CATEGORY = "Testing/Flow"
+
+    def for_loop_open(self, remaining, **kwargs):
+        graph = GraphBuilder()
+        if "initial_value0" in kwargs:
+            remaining = kwargs["initial_value0"]
+        while_open = graph.node("TestWhileLoopOpen", condition=remaining, initial_value0=remaining, **{(f"initial_value{i}"): kwargs.get(f"initial_value{i}", None) for i in range(1, NUM_FLOW_SOCKETS)})
+        outputs = [kwargs.get(f"initial_value{i}", None) for i in range(1, NUM_FLOW_SOCKETS)]
+        return {
+            "result": tuple(["stub", remaining] + outputs),
+            "expand": graph.finalize(),
+        }
+
+@VariantSupport()
+class TestForLoopClose:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "flow_control": ("FLOW_CONTROL", {"rawLink": True}),
+            },
+            "optional": {
+                f"initial_value{i}": ("*",{"rawLink": True}) for i in range(1, NUM_FLOW_SOCKETS)
+            },
+        }
+
+    RETURN_TYPES = tuple(["*"] * (NUM_FLOW_SOCKETS-1))
+    RETURN_NAMES = tuple([f"value{i}" for i in range(1, NUM_FLOW_SOCKETS)])
+    FUNCTION = "for_loop_close"
+
+    CATEGORY = "Testing/Flow"
+
+    def for_loop_close(self, flow_control, **kwargs):
+        graph = GraphBuilder()
+        while_open = flow_control[0]
+        sub = graph.node("TestIntMathOperation", operation="subtract", a=[while_open,1], b=1)
+        cond = graph.node("TestToBoolNode", value=sub.out(0))
+        input_values = {f"initial_value{i}": kwargs.get(f"initial_value{i}", None) for i in range(1, NUM_FLOW_SOCKETS)}
+        while_close = graph.node("TestWhileLoopClose",
+                flow_control=flow_control,
+                condition=cond.out(0),
+                initial_value0=sub.out(0),
+                **input_values)
+        return {
+            "result": tuple([while_close.out(i) for i in range(1, NUM_FLOW_SOCKETS)]),
+            "expand": graph.finalize(),
+        }
+
+NUM_LIST_SOCKETS = 10
+@VariantSupport()
+class TestMakeListNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "value1": ("*",),
+            },
+            "optional": {
+                f"value{i}": ("*",) for i in range(1, NUM_LIST_SOCKETS)
+            },
+        }
+
+    RETURN_TYPES = ("*",)
+    FUNCTION = "make_list"
+    OUTPUT_IS_LIST = (True,)
+
+    CATEGORY = "Testing/Lists"
+
+    def make_list(self, **kwargs):
+        result = []
+        for i in range(NUM_LIST_SOCKETS):
+            if f"value{i}" in kwargs:
+                result.append(kwargs[f"value{i}"])
+        return (result,)
+
+UTILITY_NODE_CLASS_MAPPINGS = {
+    "TestAccumulateNode": TestAccumulateNode,
+    "TestAccumulationHeadNode": TestAccumulationHeadNode,
+    "TestAccumulationTailNode": TestAccumulationTailNode,
+    "TestAccumulationToListNode": TestAccumulationToListNode,
+    "TestListToAccumulationNode": TestListToAccumulationNode,
+    "TestAccumulationGetLengthNode": TestAccumulationGetLengthNode,
+    "TestAccumulationGetItemNode": TestAccumulationGetItemNode,
+    "TestAccumulationSetItemNode": TestAccumulationSetItemNode,
+    "TestForLoopOpen": TestForLoopOpen,
+    "TestForLoopClose": TestForLoopClose,
+    "TestIntMathOperation": TestIntMathOperation,
+    "TestMakeListNode": TestMakeListNode,
+}
+UTILITY_NODE_DISPLAY_NAME_MAPPINGS = {
+    "TestAccumulateNode": "Accumulate",
+    "TestAccumulationHeadNode": "Accumulation Head",
+    "TestAccumulationTailNode": "Accumulation Tail",
+    "TestAccumulationToListNode": "Accumulation to List",
+    "TestListToAccumulationNode": "List to Accumulation",
+    "TestAccumulationGetLengthNode": "Accumulation Get Length",
+    "TestAccumulationGetItemNode": "Accumulation Get Item",
+    "TestAccumulationSetItemNode": "Accumulation Set Item",
+    "TestForLoopOpen": "For Loop Open",
+    "TestForLoopClose": "For Loop Close",
+    "TestIntMathOperation": "Int Math Operation",
+    "TestMakeListNode": "Make List",
+}

web/assets/userSelection-BGzn1LuN.css

@@ -50,6 +50,11 @@
     margin-top: 10px;
 }
 
+.comfy-user-selection input::-moz-placeholder {
+    color: var(--descrip-text);
+    opacity: 1;
+}
+
 .comfy-user-selection input::placeholder {
     color: var(--descrip-text);
     opacity: 1;
@@ -67,12 +72,8 @@
     margin: 10px 0;
     padding: 10px;
     display: block;
-    text-align: center;
     width: 100%;
     color: var(--descrip-text);
-}
-
-.comfy-user-selection-inner .or-separator {
     overflow: hidden;
     text-align: center;
     margin-left: -10px;

web/assets/userSelection-BifVfRyx.js

@@ -0,0 +1,138 @@
+var __async = (__this, __arguments, generator) => {
+  return new Promise((resolve, reject) => {
+    var fulfilled = (value) => {
+      try {
+        step(generator.next(value));
+      } catch (e) {
+        reject(e);
+      }
+    };
+    var rejected = (value) => {
+      try {
+        step(generator.throw(value));
+      } catch (e) {
+        reject(e);
+      }
+    };
+    var step = (x) => x.done ? resolve(x.value) : Promise.resolve(x.value).then(fulfilled, rejected);
+    step((generator = generator.apply(__this, __arguments)).next());
+  });
+};
+import { j as createSpinner, g as api, $ as $el } from "./index-DIiqwEjy.js";
+class UserSelectionScreen {
+  show(users, user) {
+    return __async(this, null, function* () {
+      const userSelection = document.getElementById("comfy-user-selection");
+      userSelection.style.display = "";
+      return new Promise((resolve) => {
+        const input = userSelection.getElementsByTagName("input")[0];
+        const select = userSelection.getElementsByTagName("select")[0];
+        const inputSection = input.closest("section");
+        const selectSection = select.closest("section");
+        const form = userSelection.getElementsByTagName("form")[0];
+        const error = userSelection.getElementsByClassName("comfy-user-error")[0];
+        const button = userSelection.getElementsByClassName(
+          "comfy-user-button-next"
+        )[0];
+        let inputActive = null;
+        input.addEventListener("focus", () => {
+          inputSection.classList.add("selected");
+          selectSection.classList.remove("selected");
+          inputActive = true;
+        });
+        select.addEventListener("focus", () => {
+          inputSection.classList.remove("selected");
+          selectSection.classList.add("selected");
+          inputActive = false;
+          select.style.color = "";
+        });
+        select.addEventListener("blur", () => {
+          if (!select.value) {
+            select.style.color = "var(--descrip-text)";
+          }
+        });
+        form.addEventListener("submit", (e) => __async(this, null, function* () {
+          var _a, _b, _c;
+          e.preventDefault();
+          if (inputActive == null) {
+            error.textContent = "Please enter a username or select an existing user.";
+          } else if (inputActive) {
+            const username = input.value.trim();
+            if (!username) {
+              error.textContent = "Please enter a username.";
+              return;
+            }
+            input.disabled = select.disabled = // @ts-expect-error
+            input.readonly = // @ts-expect-error
+            select.readonly = true;
+            const spinner = createSpinner();
+            button.prepend(spinner);
+            try {
+              const resp = yield api.createUser(username);
+              if (resp.status >= 300) {
+                let message = "Error creating user: " + resp.status + " " + resp.statusText;
+                try {
+                  const res = yield resp.json();
+                  if (res.error) {
+                    message = res.error;
+                  }
+                } catch (error2) {
+                }
+                throw new Error(message);
+              }
+              resolve({ username, userId: yield resp.json(), created: true });
+            } catch (err) {
+              spinner.remove();
+              error.textContent = (_c = (_b = (_a = err.message) != null ? _a : err.statusText) != null ? _b : err) != null ? _c : "An unknown error occurred.";
+              input.disabled = select.disabled = // @ts-expect-error
+              input.readonly = // @ts-expect-error
+              select.readonly = false;
+              return;
+            }
+          } else if (!select.value) {
+            error.textContent = "Please select an existing user.";
+            return;
+          } else {
+            resolve({
+              username: users[select.value],
+              userId: select.value,
+              created: false
+            });
+          }
+        }));
+        if (user) {
+          const name = localStorage["Comfy.userName"];
+          if (name) {
+            input.value = name;
+          }
+        }
+        if (input.value) {
+          input.focus();
+        }
+        const userIds = Object.keys(users != null ? users : {});
+        if (userIds.length) {
+          for (const u of userIds) {
+            $el("option", { textContent: users[u], value: u, parent: select });
+          }
+          select.style.color = "var(--descrip-text)";
+          if (select.value) {
+            select.focus();
+          }
+        } else {
+          userSelection.classList.add("no-users");
+          input.focus();
+        }
+      }).then((r) => {
+        userSelection.remove();
+        return r;
+      });
+    });
+  }
+}
+window.comfyAPI = window.comfyAPI || {};
+window.comfyAPI.userSelection = window.comfyAPI.userSelection || {};
+window.comfyAPI.userSelection.UserSelectionScreen = UserSelectionScreen;
+export {
+  UserSelectionScreen
+};
+//# sourceMappingURL=userSelection-BifVfRyx.js.map

web/extensions/core/clipspace.js

@@ -1,166 +1,2 @@
-import { app } from "../../scripts/app.js";
-import { ComfyDialog, $el } from "../../scripts/ui.js";
-import { ComfyApp } from "../../scripts/app.js";
-
-export class ClipspaceDialog extends ComfyDialog {
-	static items = [];
-	static instance = null;
-
-	static registerButton(name, contextPredicate, callback) {
-		const item =
-			$el("button", {
-				type: "button",
-				textContent: name,
-				contextPredicate: contextPredicate,
-				onclick: callback
-			})
-
-		ClipspaceDialog.items.push(item);
-	}
-
-	static invalidatePreview() {
-		if(ComfyApp.clipspace && ComfyApp.clipspace.imgs && ComfyApp.clipspace.imgs.length > 0) {
-			const img_preview = document.getElementById("clipspace_preview");
-			if(img_preview) {
-				img_preview.src = ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src;
-				img_preview.style.maxHeight = "100%";
-				img_preview.style.maxWidth = "100%";
-			}
-		}
-	}
-
-	static invalidate() {
-		if(ClipspaceDialog.instance) {
-			const self = ClipspaceDialog.instance;
-			// allow reconstruct controls when copying from non-image to image content.
-			const children = $el("div.comfy-modal-content", [ self.createImgSettings(), ...self.createButtons() ]);
-
-			if(self.element) {
-				// update
-				self.element.removeChild(self.element.firstChild);
-				self.element.appendChild(children);
-			}
-			else {
-				// new
-				self.element = $el("div.comfy-modal", { parent: document.body }, [children,]);
-			}
-
-			if(self.element.children[0].children.length <= 1) {
-				self.element.children[0].appendChild($el("p", {}, ["Unable to find the features to edit content of a format stored in the current Clipspace."]));
-			}
-
-			ClipspaceDialog.invalidatePreview();
-		}
-	}
-
-	constructor() {
-		super();
-	}
-
-	createButtons(self) {
-		const buttons = [];
-
-		for(let idx in ClipspaceDialog.items) {
-			const item = ClipspaceDialog.items[idx];
-			if(!item.contextPredicate || item.contextPredicate())
-				buttons.push(ClipspaceDialog.items[idx]);
-		}
-
-		buttons.push(
-			$el("button", {
-				type: "button",
-				textContent: "Close",
-				onclick: () => { this.close(); }
-			})
-		);
-
-		return buttons;
-	}
-
-	createImgSettings() {
-		if(ComfyApp.clipspace.imgs) {
-			const combo_items = [];
-			const imgs = ComfyApp.clipspace.imgs;
-
-			for(let i=0; i < imgs.length; i++) {
-				combo_items.push($el("option", {value:i}, [`${i}`]));
-			}
-
-			const combo1 = $el("select",
-				{id:"clipspace_img_selector", onchange:(event) => {
-					ComfyApp.clipspace['selectedIndex'] = event.target.selectedIndex;
-					ClipspaceDialog.invalidatePreview();
-				} }, combo_items);
-
-			const row1 =
-				$el("tr", {},
-						[
-							$el("td", {}, [$el("font", {color:"white"}, ["Select Image"])]),
-							$el("td", {}, [combo1])
-						]);
-
-
-			const combo2 = $el("select",
-								{id:"clipspace_img_paste_mode", onchange:(event) => {
-									ComfyApp.clipspace['img_paste_mode'] = event.target.value;
-								} },
-								[
-									$el("option", {value:'selected'}, 'selected'),
-									$el("option", {value:'all'}, 'all')
-								]);
-			combo2.value = ComfyApp.clipspace['img_paste_mode'];
-
-			const row2 =
-				$el("tr", {},
-						[
-							$el("td", {}, [$el("font", {color:"white"}, ["Paste Mode"])]),
-							$el("td", {}, [combo2])
-						]);
-
-			const td = $el("td", {align:'center', width:'100px', height:'100px', colSpan:'2'},
-								[ $el("img",{id:"clipspace_preview", ondragstart:() => false},[]) ]);
-
-			const row3 =
-				$el("tr", {}, [td]);
-
-			return $el("table", {}, [row1, row2, row3]);
-		}
-		else {
-			return [];
-		}
-	}
-
-	createImgPreview() {
-		if(ComfyApp.clipspace.imgs) {
-			return $el("img",{id:"clipspace_preview", ondragstart:() => false});
-		}
-		else
-			return [];
-	}
-
-	show() {
-		const img_preview = document.getElementById("clipspace_preview");
-		ClipspaceDialog.invalidate();
-		
-		this.element.style.display = "block";
-	}
-}
-
-app.registerExtension({
-	name: "Comfy.Clipspace",
-	init(app) {
-		app.openClipspace =
-			function () {
-				if(!ClipspaceDialog.instance) {
-					ClipspaceDialog.instance = new ClipspaceDialog(app);
-					ComfyApp.clipspace_invalidate_handler = ClipspaceDialog.invalidate;
-				}
-
-				if(ComfyApp.clipspace) {
-					ClipspaceDialog.instance.show();
-				}
-				else
-					app.ui.dialog.show("Clipspace is Empty!");
-			};
-	}
-});
+// Shim for extensions\core\clipspace.ts
+export const ClipspaceDialog = window.comfyAPI.clipspace.ClipspaceDialog;

web/extensions/core/colorPalette.js

@@ -1,785 +0,0 @@
-import {app} from "../../scripts/app.js";
-import {$el} from "../../scripts/ui.js";
-
-// Manage color palettes
-
-const colorPalettes = {
-	"dark": {
-		"id": "dark",
-		"name": "Dark (Default)",
-		"colors": {
-			"node_slot": {
-				"CLIP": "#FFD500", // bright yellow
-				"CLIP_VISION": "#A8DADC", // light blue-gray
-				"CLIP_VISION_OUTPUT": "#ad7452", // rusty brown-orange
-				"CONDITIONING": "#FFA931", // vibrant orange-yellow
-				"CONTROL_NET": "#6EE7B7", // soft mint green
-				"IMAGE": "#64B5F6", // bright sky blue
-				"LATENT": "#FF9CF9", // light pink-purple
-				"MASK": "#81C784", // muted green
-				"MODEL": "#B39DDB", // light lavender-purple
-				"STYLE_MODEL": "#C2FFAE", // light green-yellow
-				"VAE": "#FF6E6E", // bright red
-				"NOISE": "#B0B0B0", // gray
-				"GUIDER": "#66FFFF", // cyan
-				"SAMPLER": "#ECB4B4", // very soft red
-				"SIGMAS": "#CDFFCD", // soft lime green
-				"TAESD": "#DCC274", // cheesecake
-			},
-			"litegraph_base": {
-				"BACKGROUND_IMAGE": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABkCAIAAAD/gAIDAAAAGXRFWHRTb2Z0d2FyZQBBZG9iZSBJbWFnZVJlYWR5ccllPAAAAQBJREFUeNrs1rEKwjAUhlETUkj3vP9rdmr1Ysammk2w5wdxuLgcMHyptfawuZX4pJSWZTnfnu/lnIe/jNNxHHGNn//HNbbv+4dr6V+11uF527arU7+u63qfa/bnmh8sWLBgwYJlqRf8MEptXPBXJXa37BSl3ixYsGDBMliwFLyCV/DeLIMFCxYsWLBMwSt4Be/NggXLYMGCBUvBK3iNruC9WbBgwYJlsGApeAWv4L1ZBgsWLFiwYJmCV/AK3psFC5bBggULloJX8BpdwXuzYMGCBctgwVLwCl7Be7MMFixYsGDBsu8FH1FaSmExVfAxBa/gvVmwYMGCZbBg/W4vAQYA5tRF9QYlv/QAAAAASUVORK5CYII=",
-				"CLEAR_BACKGROUND_COLOR": "#222",
-				"NODE_TITLE_COLOR": "#999",
-				"NODE_SELECTED_TITLE_COLOR": "#FFF",
-				"NODE_TEXT_SIZE": 14,
-				"NODE_TEXT_COLOR": "#AAA",
-				"NODE_SUBTEXT_SIZE": 12,
-				"NODE_DEFAULT_COLOR": "#333",
-				"NODE_DEFAULT_BGCOLOR": "#353535",
-				"NODE_DEFAULT_BOXCOLOR": "#666",
-				"NODE_DEFAULT_SHAPE": "box",
-				"NODE_BOX_OUTLINE_COLOR": "#FFF",
-				"DEFAULT_SHADOW_COLOR": "rgba(0,0,0,0.5)",
-				"DEFAULT_GROUP_FONT": 24,
-
-				"WIDGET_BGCOLOR": "#222",
-				"WIDGET_OUTLINE_COLOR": "#666",
-				"WIDGET_TEXT_COLOR": "#DDD",
-				"WIDGET_SECONDARY_TEXT_COLOR": "#999",
-
-				"LINK_COLOR": "#9A9",
-				"EVENT_LINK_COLOR": "#A86",
-				"CONNECTING_LINK_COLOR": "#AFA",
-			},
-			"comfy_base": {
-				"fg-color": "#fff",
-				"bg-color": "#202020",
-				"comfy-menu-bg": "#353535",
-				"comfy-input-bg": "#222",
-				"input-text": "#ddd",
-				"descrip-text": "#999",
-				"drag-text": "#ccc",
-				"error-text": "#ff4444",
-				"border-color": "#4e4e4e",
-				"tr-even-bg-color": "#222",
-				"tr-odd-bg-color": "#353535",
-				"content-bg": "#4e4e4e",
-				"content-fg": "#fff",
-				"content-hover-bg": "#222",
-				"content-hover-fg": "#fff"
-			}
-		},
-	},
-	"light": {
-		"id": "light",
-		"name": "Light",
-		"colors": {
-			"node_slot": {
-				"CLIP": "#FFA726", // orange
-				"CLIP_VISION": "#5C6BC0", // indigo
-				"CLIP_VISION_OUTPUT": "#8D6E63", // brown
-				"CONDITIONING": "#EF5350", // red
-				"CONTROL_NET": "#66BB6A", // green
-				"IMAGE": "#42A5F5", // blue
-				"LATENT": "#AB47BC", // purple
-				"MASK": "#9CCC65", // light green
-				"MODEL": "#7E57C2", // deep purple
-				"STYLE_MODEL": "#D4E157", // lime
-				"VAE": "#FF7043", // deep orange
-			},
-			"litegraph_base": {
-				"BACKGROUND_IMAGE": "data:image/gif;base64,R0lGODlhZABkALMAAAAAAP///+vr6+rq6ujo6Ofn5+bm5uXl5d3d3f///wAAAAAAAAAAAAAAAAAAAAAAACH5BAEAAAkALAAAAABkAGQAAAT/UMhJq7046827HkcoHkYxjgZhnGG6si5LqnIM0/fL4qwwIMAg0CAsEovBIxKhRDaNy2GUOX0KfVFrssrNdpdaqTeKBX+dZ+jYvEaTf+y4W66mC8PUdrE879f9d2mBeoNLfH+IhYBbhIx2jkiHiomQlGKPl4uZe3CaeZifnnijgkESBqipqqusra6vsLGys62SlZO4t7qbuby7CLa+wqGWxL3Gv3jByMOkjc2lw8vOoNSi0czAncXW3Njdx9Pf48/Z4Kbbx+fQ5evZ4u3k1fKR6cn03vHlp7T9/v8A/8Gbp4+gwXoFryXMB2qgwoMMHyKEqA5fxX322FG8tzBcRnMW/zlulPbRncmQGidKjMjyYsOSKEF2FBlJQMCbOHP6c9iSZs+UnGYCdbnSo1CZI5F64kn0p1KnTH02nSoV3dGTV7FFHVqVq1dtWcMmVQZTbNGu72zqXMuW7danVL+6e4t1bEy6MeueBYLXrNO5Ze36jQtWsOG97wIj1vt3St/DjTEORss4nNq2mDP3e7w4r1bFkSET5hy6s2TRlD2/mSxXtSHQhCunXo26NevCpmvD/UU6tuullzULH76q92zdZG/Ltv1a+W+osI/nRmyc+fRi1Xdbh+68+0vv10dH3+77KD/i6IdnX669/frn5Zsjh4/2PXju8+8bzc9/6fj27LFnX11/+IUnXWl7BJfegm79FyB9JOl3oHgSklefgxAC+FmFGpqHIYcCfkhgfCohSKKJVo044YUMttggiBkmp6KFXw1oII24oYhjiDByaKOOHcp3Y5BD/njikSkO+eBREQAAOw==",
-				"CLEAR_BACKGROUND_COLOR": "lightgray",
-				"NODE_TITLE_COLOR": "#222",
-				"NODE_SELECTED_TITLE_COLOR": "#000",
-				"NODE_TEXT_SIZE": 14,
-				"NODE_TEXT_COLOR": "#444",
-				"NODE_SUBTEXT_SIZE": 12,
-				"NODE_DEFAULT_COLOR": "#F7F7F7",
-				"NODE_DEFAULT_BGCOLOR": "#F5F5F5",
-				"NODE_DEFAULT_BOXCOLOR": "#CCC",
-				"NODE_DEFAULT_SHAPE": "box",
-				"NODE_BOX_OUTLINE_COLOR": "#000",
-				"DEFAULT_SHADOW_COLOR": "rgba(0,0,0,0.1)",
-				"DEFAULT_GROUP_FONT": 24,
-
-				"WIDGET_BGCOLOR": "#D4D4D4",
-				"WIDGET_OUTLINE_COLOR": "#999",
-				"WIDGET_TEXT_COLOR": "#222",
-				"WIDGET_SECONDARY_TEXT_COLOR": "#555",
-
-				"LINK_COLOR": "#4CAF50",
-				"EVENT_LINK_COLOR": "#FF9800",
-				"CONNECTING_LINK_COLOR": "#2196F3",
-			},
-			"comfy_base": {
-				"fg-color": "#222",
-				"bg-color": "#DDD",
-				"comfy-menu-bg": "#F5F5F5",
-				"comfy-input-bg": "#C9C9C9",
-				"input-text": "#222",
-				"descrip-text": "#444",
-				"drag-text": "#555",
-				"error-text": "#F44336",
-				"border-color": "#888",
-				"tr-even-bg-color": "#f9f9f9",
-				"tr-odd-bg-color": "#fff",
-				"content-bg": "#e0e0e0",
-				"content-fg": "#222",
-				"content-hover-bg": "#adadad",
-				"content-hover-fg": "#222"
-			}
-		},
-	},
-	"solarized": {
-		"id": "solarized",
-		"name": "Solarized",
-		"colors": {
-			"node_slot": {
-				"CLIP": "#2AB7CA", // light blue
-				"CLIP_VISION": "#6c71c4", // blue violet
-				"CLIP_VISION_OUTPUT": "#859900", // olive green
-				"CONDITIONING": "#d33682", // magenta
-				"CONTROL_NET": "#d1ffd7", // light mint green
-				"IMAGE": "#5940bb", // deep blue violet
-				"LATENT": "#268bd2", // blue
-				"MASK": "#CCC9E7", // light purple-gray
-				"MODEL": "#dc322f", // red
-				"STYLE_MODEL": "#1a998a", // teal
-				"UPSCALE_MODEL": "#054A29", // dark green
-				"VAE": "#facfad", // light pink-orange
-			},
-			"litegraph_base": {
-				"NODE_TITLE_COLOR": "#fdf6e3", // Base3
-				"NODE_SELECTED_TITLE_COLOR": "#A9D400",
-				"NODE_TEXT_SIZE": 14,
-				"NODE_TEXT_COLOR": "#657b83", // Base00
-				"NODE_SUBTEXT_SIZE": 12,
-				"NODE_DEFAULT_COLOR": "#094656",
-				"NODE_DEFAULT_BGCOLOR": "#073642", // Base02
-				"NODE_DEFAULT_BOXCOLOR": "#839496", // Base0
-				"NODE_DEFAULT_SHAPE": "box",
-				"NODE_BOX_OUTLINE_COLOR": "#fdf6e3", // Base3
-				"DEFAULT_SHADOW_COLOR": "rgba(0,0,0,0.5)",
-				"DEFAULT_GROUP_FONT": 24,
-
-				"WIDGET_BGCOLOR": "#002b36", // Base03
-				"WIDGET_OUTLINE_COLOR": "#839496", // Base0
-				"WIDGET_TEXT_COLOR": "#fdf6e3", // Base3
-				"WIDGET_SECONDARY_TEXT_COLOR": "#93a1a1", // Base1
-
-				"LINK_COLOR": "#2aa198", // Solarized Cyan
-				"EVENT_LINK_COLOR": "#268bd2", // Solarized Blue
-				"CONNECTING_LINK_COLOR": "#859900", // Solarized Green
-			},
-			"comfy_base": {
-				"fg-color": "#fdf6e3", // Base3
-				"bg-color": "#002b36", // Base03
-				"comfy-menu-bg": "#073642", // Base02
-				"comfy-input-bg": "#002b36", // Base03
-				"input-text": "#93a1a1", // Base1
-				"descrip-text": "#586e75", // Base01
-				"drag-text": "#839496", // Base0
-				"error-text": "#dc322f", // Solarized Red
-				"border-color": "#657b83", // Base00
-				"tr-even-bg-color": "#002b36",
-				"tr-odd-bg-color": "#073642",
-				"content-bg": "#657b83",
-				"content-fg": "#fdf6e3",
-				"content-hover-bg": "#002b36",
-				"content-hover-fg": "#fdf6e3"
-			}
-		},
-	},
-	"arc": {
-		"id": "arc",
-		"name": "Arc",
-		"colors": {
-			"node_slot": {
-			  "BOOLEAN": "",
-			  "CLIP": "#eacb8b",
-			  "CLIP_VISION": "#A8DADC",
-			  "CLIP_VISION_OUTPUT": "#ad7452",
-			  "CONDITIONING": "#cf876f",
-			  "CONTROL_NET": "#00d78d",
-			  "CONTROL_NET_WEIGHTS": "",
-			  "FLOAT": "",
-			  "GLIGEN": "",
-			  "IMAGE": "#80a1c0",
-			  "IMAGEUPLOAD": "",
-			  "INT": "",
-			  "LATENT": "#b38ead",
-			  "LATENT_KEYFRAME": "",
-			  "MASK": "#a3bd8d",
-			  "MODEL": "#8978a7",
-			  "SAMPLER": "",
-			  "SIGMAS": "",
-			  "STRING": "",
-			  "STYLE_MODEL": "#C2FFAE",
-			  "T2I_ADAPTER_WEIGHTS": "",
-			  "TAESD": "#DCC274",
-			  "TIMESTEP_KEYFRAME": "",
-			  "UPSCALE_MODEL": "",
-			  "VAE": "#be616b"
-			},
-			"litegraph_base": {
-			  "BACKGROUND_IMAGE": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABkCAYAAABw4pVUAAAACXBIWXMAAAsTAAALEwEAmpwYAAABcklEQVR4nO3YMUoDARgF4RfxBqZI6/0vZqFn0MYtrLIQMFN8U6V4LAtD+Jm9XG/v30OGl2e/AP7yevz4+vx45nvgF/+QGITEICQGITEIiUFIjNNC3q43u3/YnRJyPOzeQ+0e220nhRzReC8e7R7bbdvl+Jal1Bs46jEIiUFIDEJiEBKDkBhKPbZT6qHdptRTu02p53DUYxASg5AYhMQgJAYhMZR6bKfUQ7tNqad2m1LP4ajHICQGITEIiUFIDEJiKPXYTqmHdptST+02pZ7DUY9BSAxCYhASg5AYhMRQ6rGdUg/tNqWe2m1KPYejHoOQGITEICQGITEIiaHUYzulHtptSj2125R6Dkc9BiExCIlBSAxCYhASQ6nHdko9tNuUemq3KfUcjnoMQmIQEoOQGITEICSGUo/tlHpotyn11G5T6jkc9RiExCAkBiExCIlBSAylHtsp9dBuU+qp3abUczjqMQiJQUgMQmIQEoOQGITE+AHFISNQrFTGuwAAAABJRU5ErkJggg==",
-			  "CLEAR_BACKGROUND_COLOR": "#2b2f38",
-			  "NODE_TITLE_COLOR": "#b2b7bd",
-			  "NODE_SELECTED_TITLE_COLOR": "#FFF",
-			  "NODE_TEXT_SIZE": 14,
-			  "NODE_TEXT_COLOR": "#AAA",
-			  "NODE_SUBTEXT_SIZE": 12,
-			  "NODE_DEFAULT_COLOR": "#2b2f38",
-			  "NODE_DEFAULT_BGCOLOR": "#242730",
-			  "NODE_DEFAULT_BOXCOLOR": "#6e7581",
-			  "NODE_DEFAULT_SHAPE": "box",
-			  "NODE_BOX_OUTLINE_COLOR": "#FFF",
-			  "DEFAULT_SHADOW_COLOR": "rgba(0,0,0,0.5)",
-			  "DEFAULT_GROUP_FONT": 22,
-			  "WIDGET_BGCOLOR": "#2b2f38",
-			  "WIDGET_OUTLINE_COLOR": "#6e7581",
-			  "WIDGET_TEXT_COLOR": "#DDD",
-			  "WIDGET_SECONDARY_TEXT_COLOR": "#b2b7bd",
-			  "LINK_COLOR": "#9A9",
-			  "EVENT_LINK_COLOR": "#A86",
-			  "CONNECTING_LINK_COLOR": "#AFA"
-			},
-			"comfy_base": {
-			  "fg-color": "#fff",
-			  "bg-color": "#2b2f38",
-			  "comfy-menu-bg": "#242730",
-			  "comfy-input-bg": "#2b2f38",
-			  "input-text": "#ddd",
-			  "descrip-text": "#b2b7bd",
-			  "drag-text": "#ccc",
-			  "error-text": "#ff4444",
-			  "border-color": "#6e7581",
-			  "tr-even-bg-color": "#2b2f38",
-			  "tr-odd-bg-color": "#242730",
-			  "content-bg": "#6e7581",
-			  "content-fg": "#fff",
-			  "content-hover-bg": "#2b2f38",
-			  "content-hover-fg": "#fff"
-			}
-		},
-	},
-	"nord": {
-		"id": "nord",
-		"name": "Nord",
-		"colors": {
-			"node_slot": {
-			  "BOOLEAN": "",
-			  "CLIP": "#eacb8b",
-			  "CLIP_VISION": "#A8DADC",
-			  "CLIP_VISION_OUTPUT": "#ad7452",
-			  "CONDITIONING": "#cf876f",
-			  "CONTROL_NET": "#00d78d",
-			  "CONTROL_NET_WEIGHTS": "",
-			  "FLOAT": "",
-			  "GLIGEN": "",
-			  "IMAGE": "#80a1c0",
-			  "IMAGEUPLOAD": "",
-			  "INT": "",
-			  "LATENT": "#b38ead",
-			  "LATENT_KEYFRAME": "",
-			  "MASK": "#a3bd8d",
-			  "MODEL": "#8978a7",
-			  "SAMPLER": "",
-			  "SIGMAS": "",
-			  "STRING": "",
-			  "STYLE_MODEL": "#C2FFAE",
-			  "T2I_ADAPTER_WEIGHTS": "",
-			  "TAESD": "#DCC274",
-			  "TIMESTEP_KEYFRAME": "",
-			  "UPSCALE_MODEL": "",
-			  "VAE": "#be616b"
-			},
-			"litegraph_base": {
-			  "BACKGROUND_IMAGE": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABkCAIAAAD/gAIDAAAACXBIWXMAAAsTAAALEwEAmpwYAAAFu2lUWHRYTUw6Y29tLmFkb2JlLnhtcAAAAAAAPD94cGFja2V0IGJlZ2luPSLvu78iIGlkPSJXNU0wTXBDZWhpSHpyZVN6TlRjemtjOWQiPz4gPHg6eG1wbWV0YSB4bWxuczp4PSJhZG9iZTpuczptZXRhLyIgeDp4bXB0az0iQWRvYmUgWE1QIENvcmUgOS4xLWMwMDEgNzkuMTQ2Mjg5OSwgMjAyMy8wNi8yNS0yMDowMTo1NSAgICAgICAgIj4gPHJkZjpSREYgeG1sbnM6cmRmPSJodHRwOi8vd3d3LnczLm9yZy8xOTk5LzAyLzIyLXJkZi1zeW50YXgtbnMjIj4gPHJkZjpEZXNjcmlwdGlvbiByZGY6YWJvdXQ9IiIgeG1sbnM6eG1wPSJodHRwOi8vbnMuYWRvYmUuY29tL3hhcC8xLjAvIiB4bWxuczpkYz0iaHR0cDovL3B1cmwub3JnL2RjL2VsZW1lbnRzLzEuMS8iIHhtbG5zOnBob3Rvc2hvcD0iaHR0cDovL25zLmFkb2JlLmNvbS9waG90b3Nob3AvMS4wLyIgeG1sbnM6eG1wTU09Imh0dHA6Ly9ucy5hZG9iZS5jb20veGFwLzEuMC9tbS8iIHhtbG5zOnN0RXZ0PSJodHRwOi8vbnMuYWRvYmUuY29tL3hhcC8xLjAvc1R5cGUvUmVzb3VyY2VFdmVudCMiIHhtcDpDcmVhdG9yVG9vbD0iQWRvYmUgUGhvdG9zaG9wIDI1LjEgKFdpbmRvd3MpIiB4bXA6Q3JlYXRlRGF0ZT0iMjAyMy0xMS0xM1QwMDoxODowMiswMTowMCIgeG1wOk1vZGlmeURhdGU9IjIwMjMtMTEtMTVUMDE6MjA6NDUrMDE6MDAiIHhtcDpNZXRhZGF0YURhdGU9IjIwMjMtMTEtMTVUMDE6MjA6NDUrMDE6MDAiIGRjOmZvcm1hdD0iaW1hZ2UvcG5nIiBwaG90b3Nob3A6Q29sb3JNb2RlPSIzIiB4bXBNTTpJbnN0YW5jZUlEPSJ4bXAuaWlkOjUwNDFhMmZjLTEzNzQtMTk0ZC1hZWY4LTYxMzM1MTVmNjUwMCIgeG1wTU06RG9jdW1lbnRJRD0ieG1wLmRpZDoyMzFiMTBiMC1iNGZiLTAyNGUtYjEyZS0zMDUzMDNjZDA3YzgiIHhtcE1NOk9yaWdpbmFsRG9jdW1lbnRJRD0ieG1wLmRpZDoyMzFiMTBiMC1iNGZiLTAyNGUtYjEyZS0zMDUzMDNjZDA3YzgiPiA8eG1wTU06SGlzdG9yeT4gPHJkZjpTZXE+IDxyZGY6bGkgc3RFdnQ6YWN0aW9uPSJjcmVhdGVkIiBzdEV2dDppbnN0YW5jZUlEPSJ4bXAuaWlkOjIzMWIxMGIwLWI0ZmItMDI0ZS1iMTJlLTMwNTMwM2NkMDdjOCIgc3RFdnQ6d2hlbj0iMjAyMy0xMS0xM1QwMDoxODowMiswMTowMCIgc3RFdnQ6c29mdHdhcmVBZ2VudD0iQWRvYmUgUGhvdG9zaG9wIDI1LjEgKFdpbmRvd3MpIi8+IDxyZGY6bGkgc3RFdnQ6YWN0aW9uPSJzYXZlZCIgc3RFdnQ6aW5zdGFuY2VJRD0ieG1wLmlpZDo1MDQxYTJmYy0xMzc0LTE5NGQtYWVmOC02MTMzNTE1ZjY1MDAiIHN0RXZ0OndoZW49IjIwMjMtMTEtMTVUMDE6MjA6NDUrMDE6MDAiIHN0RXZ0OnNvZnR3YXJlQWdlbnQ9IkFkb2JlIFBob3Rvc2hvcCAyNS4xIChXaW5kb3dzKSIgc3RFdnQ6Y2hhbmdlZD0iLyIvPiA8L3JkZjpTZXE+IDwveG1wTU06SGlzdG9yeT4gPC9yZGY6RGVzY3JpcHRpb24+IDwvcmRmOlJERj4gPC94OnhtcG1ldGE+IDw/eHBhY2tldCBlbmQ9InIiPz73jWg/AAAAyUlEQVR42u3WKwoAIBRFQRdiMb1idv9Lsxn9gEFw4Dbb8JCTojbbXEJwjJVL2HKwYMGCBQuWLbDmjr+9zrBGjHl1WVcvy2DBggULFizTWQpewSt4HzwsgwULFiwFr7MUvMtS8D54WLBgGSxYCl7BK3iXZbBgwYIFC5bpLAWv4BW8Dx6WwYIFC5aC11kK3mUpeB88LFiwDBYsBa/gFbzLMliwYMGCBct0loJX8AreBw/LYMGCBUvB6ywF77IUvA8eFixYBgsWrNfWAZPltufdad+1AAAAAElFTkSuQmCC",
-			  "CLEAR_BACKGROUND_COLOR": "#212732",
-			  "NODE_TITLE_COLOR": "#999",
-			  "NODE_SELECTED_TITLE_COLOR": "#e5eaf0",
-			  "NODE_TEXT_SIZE": 14,
-			  "NODE_TEXT_COLOR": "#bcc2c8",
-			  "NODE_SUBTEXT_SIZE": 12,
-			  "NODE_DEFAULT_COLOR": "#2e3440",
-			  "NODE_DEFAULT_BGCOLOR": "#161b22",
-			  "NODE_DEFAULT_BOXCOLOR": "#545d70",
-			  "NODE_DEFAULT_SHAPE": "box",
-			  "NODE_BOX_OUTLINE_COLOR": "#e5eaf0",
-			  "DEFAULT_SHADOW_COLOR": "rgba(0,0,0,0.5)",
-			  "DEFAULT_GROUP_FONT": 24,
-			  "WIDGET_BGCOLOR": "#2e3440",
-			  "WIDGET_OUTLINE_COLOR": "#545d70",
-			  "WIDGET_TEXT_COLOR": "#bcc2c8",
-			  "WIDGET_SECONDARY_TEXT_COLOR": "#999",
-			  "LINK_COLOR": "#9A9",
-			  "EVENT_LINK_COLOR": "#A86",
-			  "CONNECTING_LINK_COLOR": "#AFA"
-			},
-			"comfy_base": {
-			  "fg-color": "#e5eaf0",
-			  "bg-color": "#2e3440",
-			  "comfy-menu-bg": "#161b22",
-			  "comfy-input-bg": "#2e3440",
-			  "input-text": "#bcc2c8",
-			  "descrip-text": "#999",
-			  "drag-text": "#ccc",
-			  "error-text": "#ff4444",
-			  "border-color": "#545d70",
-			  "tr-even-bg-color": "#2e3440",
-			  "tr-odd-bg-color": "#161b22",
-			  "content-bg": "#545d70",
-			  "content-fg": "#e5eaf0",
-			  "content-hover-bg": "#2e3440",
-			  "content-hover-fg": "#e5eaf0"
-			}
-		},
-	},
-	"github": {
-		"id": "github",
-		"name": "Github",
-		"colors": {
-			"node_slot": {
-			  "BOOLEAN": "",
-			  "CLIP": "#eacb8b",
-			  "CLIP_VISION": "#A8DADC",
-			  "CLIP_VISION_OUTPUT": "#ad7452",
-			  "CONDITIONING": "#cf876f",
-			  "CONTROL_NET": "#00d78d",
-			  "CONTROL_NET_WEIGHTS": "",
-			  "FLOAT": "",
-			  "GLIGEN": "",
-			  "IMAGE": "#80a1c0",
-			  "IMAGEUPLOAD": "",
-			  "INT": "",
-			  "LATENT": "#b38ead",
-			  "LATENT_KEYFRAME": "",
-			  "MASK": "#a3bd8d",
-			  "MODEL": "#8978a7",
-			  "SAMPLER": "",
-			  "SIGMAS": "",
-			  "STRING": "",
-			  "STYLE_MODEL": "#C2FFAE",
-			  "T2I_ADAPTER_WEIGHTS": "",
-			  "TAESD": "#DCC274",
-			  "TIMESTEP_KEYFRAME": "",
-			  "UPSCALE_MODEL": "",
-			  "VAE": "#be616b"
-			},
-			"litegraph_base": {
-			  "BACKGROUND_IMAGE": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABkCAIAAAD/gAIDAAAACXBIWXMAAAsTAAALEwEAmpwYAAAGlmlUWHRYTUw6Y29tLmFkb2JlLnhtcAAAAAAAPD94cGFja2V0IGJlZ2luPSLvu78iIGlkPSJXNU0wTXBDZWhpSHpyZVN6TlRjemtjOWQiPz4gPHg6eG1wbWV0YSB4bWxuczp4PSJhZG9iZTpuczptZXRhLyIgeDp4bXB0az0iQWRvYmUgWE1QIENvcmUgOS4xLWMwMDEgNzkuMTQ2Mjg5OSwgMjAyMy8wNi8yNS0yMDowMTo1NSAgICAgICAgIj4gPHJkZjpSREYgeG1sbnM6cmRmPSJodHRwOi8vd3d3LnczLm9yZy8xOTk5LzAyLzIyLXJkZi1zeW50YXgtbnMjIj4gPHJkZjpEZXNjcmlwdGlvbiByZGY6YWJvdXQ9IiIgeG1sbnM6eG1wPSJodHRwOi8vbnMuYWRvYmUuY29tL3hhcC8xLjAvIiB4bWxuczpkYz0iaHR0cDovL3B1cmwub3JnL2RjL2VsZW1lbnRzLzEuMS8iIHhtbG5zOnBob3Rvc2hvcD0iaHR0cDovL25zLmFkb2JlLmNvbS9waG90b3Nob3AvMS4wLyIgeG1sbnM6eG1wTU09Imh0dHA6Ly9ucy5hZG9iZS5jb20veGFwLzEuMC9tbS8iIHhtbG5zOnN0RXZ0PSJodHRwOi8vbnMuYWRvYmUuY29tL3hhcC8xLjAvc1R5cGUvUmVzb3VyY2VFdmVudCMiIHhtcDpDcmVhdG9yVG9vbD0iQWRvYmUgUGhvdG9zaG9wIDI1LjEgKFdpbmRvd3MpIiB4bXA6Q3JlYXRlRGF0ZT0iMjAyMy0xMS0xM1QwMDoxODowMiswMTowMCIgeG1wOk1vZGlmeURhdGU9IjIwMjMtMTEtMTVUMDI6MDQ6NTkrMDE6MDAiIHhtcDpNZXRhZGF0YURhdGU9IjIwMjMtMTEtMTVUMDI6MDQ6NTkrMDE6MDAiIGRjOmZvcm1hdD0iaW1hZ2UvcG5nIiBwaG90b3Nob3A6Q29sb3JNb2RlPSIzIiB4bXBNTTpJbnN0YW5jZUlEPSJ4bXAuaWlkOmIyYzRhNjA5LWJmYTctYTg0MC1iOGFlLTk3MzE2ZjM1ZGIyNyIgeG1wTU06RG9jdW1lbnRJRD0iYWRvYmU6ZG9jaWQ6cGhvdG9zaG9wOjk0ZmNlZGU4LTE1MTctZmQ0MC04ZGU3LWYzOTgxM2E3ODk5ZiIgeG1wTU06T3JpZ2luYWxEb2N1bWVudElEPSJ4bXAuZGlkOjIzMWIxMGIwLWI0ZmItMDI0ZS1iMTJlLTMwNTMwM2NkMDdjOCI+IDx4bXBNTTpIaXN0b3J5PiA8cmRmOlNlcT4gPHJkZjpsaSBzdEV2dDphY3Rpb249ImNyZWF0ZWQiIHN0RXZ0Omluc3RhbmNlSUQ9InhtcC5paWQ6MjMxYjEwYjAtYjRmYi0wMjRlLWIxMmUtMzA1MzAzY2QwN2M4IiBzdEV2dDp3aGVuPSIyMDIzLTExLTEzVDAwOjE4OjAyKzAxOjAwIiBzdEV2dDpzb2Z0d2FyZUFnZW50PSJBZG9iZSBQaG90b3Nob3AgMjUuMSAoV2luZG93cykiLz4gPHJkZjpsaSBzdEV2dDphY3Rpb249InNhdmVkIiBzdEV2dDppbnN0YW5jZUlEPSJ4bXAuaWlkOjQ4OWY1NzlmLTJkNjUtZWQ0Zi04OTg0LTA4NGE2MGE1ZTMzNSIgc3RFdnQ6d2hlbj0iMjAyMy0xMS0xNVQwMjowNDo1OSswMTowMCIgc3RFdnQ6c29mdHdhcmVBZ2VudD0iQWRvYmUgUGhvdG9zaG9wIDI1LjEgKFdpbmRvd3MpIiBzdEV2dDpjaGFuZ2VkPSIvIi8+IDxyZGY6bGkgc3RFdnQ6YWN0aW9uPSJzYXZlZCIgc3RFdnQ6aW5zdGFuY2VJRD0ieG1wLmlpZDpiMmM0YTYwOS1iZmE3LWE4NDAtYjhhZS05NzMxNmYzNWRiMjciIHN0RXZ0OndoZW49IjIwMjMtMTEtMTVUMDI6MDQ6NTkrMDE6MDAiIHN0RXZ0OnNvZnR3YXJlQWdlbnQ9IkFkb2JlIFBob3Rvc2hvcCAyNS4xIChXaW5kb3dzKSIgc3RFdnQ6Y2hhbmdlZD0iLyIvPiA8L3JkZjpTZXE+IDwveG1wTU06SGlzdG9yeT4gPC9yZGY6RGVzY3JpcHRpb24+IDwvcmRmOlJERj4gPC94OnhtcG1ldGE+IDw/eHBhY2tldCBlbmQ9InIiPz4OTe6GAAAAx0lEQVR42u3WMQoAIQxFwRzJys77X8vSLiRgITif7bYbgrwYc/mKXyBoY4VVBgsWLFiwYFmOlTv+9jfDOjHmr8u6eVkGCxYsWLBgmc5S8ApewXvgYRksWLBgKXidpeBdloL3wMOCBctgwVLwCl7BuyyDBQsWLFiwTGcpeAWv4D3wsAwWLFiwFLzOUvAuS8F74GHBgmWwYCl4Ba/gXZbBggULFixYprMUvIJX8B54WAYLFixYCl5nKXiXpeA98LBgwTJYsGC9tg1o8f4TTtqzNQAAAABJRU5ErkJggg==",
-			  "CLEAR_BACKGROUND_COLOR": "#040506",
-			  "NODE_TITLE_COLOR": "#999",
-			  "NODE_SELECTED_TITLE_COLOR": "#e5eaf0",
-			  "NODE_TEXT_SIZE": 14,
-			  "NODE_TEXT_COLOR": "#bcc2c8",
-			  "NODE_SUBTEXT_SIZE": 12,
-			  "NODE_DEFAULT_COLOR": "#161b22",
-			  "NODE_DEFAULT_BGCOLOR": "#13171d",
-			  "NODE_DEFAULT_BOXCOLOR": "#30363d",
-			  "NODE_DEFAULT_SHAPE": "box",
-			  "NODE_BOX_OUTLINE_COLOR": "#e5eaf0",
-			  "DEFAULT_SHADOW_COLOR": "rgba(0,0,0,0.5)",
-			  "DEFAULT_GROUP_FONT": 24,
-			  "WIDGET_BGCOLOR": "#161b22",
-			  "WIDGET_OUTLINE_COLOR": "#30363d",
-			  "WIDGET_TEXT_COLOR": "#bcc2c8",
-			  "WIDGET_SECONDARY_TEXT_COLOR": "#999",
-			  "LINK_COLOR": "#9A9",
-			  "EVENT_LINK_COLOR": "#A86",
-			  "CONNECTING_LINK_COLOR": "#AFA"
-			},
-			"comfy_base": {
-			  "fg-color": "#e5eaf0",
-			  "bg-color": "#161b22",
-			  "comfy-menu-bg": "#13171d",
-			  "comfy-input-bg": "#161b22",
-			  "input-text": "#bcc2c8",
-			  "descrip-text": "#999",
-			  "drag-text": "#ccc",
-			  "error-text": "#ff4444",
-			  "border-color": "#30363d",
-			  "tr-even-bg-color": "#161b22",
-			  "tr-odd-bg-color": "#13171d",
-			  "content-bg": "#30363d",
-			  "content-fg": "#e5eaf0",
-			  "content-hover-bg": "#161b22",
-			  "content-hover-fg": "#e5eaf0"
-			}
-		},
-	}
-};
-
-const id = "Comfy.ColorPalette";
-const idCustomColorPalettes = "Comfy.CustomColorPalettes";
-const defaultColorPaletteId = "dark";
-const els = {}
-// const ctxMenu = LiteGraph.ContextMenu;
-app.registerExtension({
-	name: id,
-	addCustomNodeDefs(node_defs) {
-		const sortObjectKeys = (unordered) => {
-			return Object.keys(unordered).sort().reduce((obj, key) => {
-				obj[key] = unordered[key];
-				return obj;
-			}, {});
-		};
-
-		function getSlotTypes() {
-			var types = [];
-
-			const defs = node_defs;
-			for (const nodeId in defs) {
-				const nodeData = defs[nodeId];
-
-				var inputs = nodeData["input"]["required"];
-				if (nodeData["input"]["optional"] !== undefined) {
-					inputs = Object.assign({}, nodeData["input"]["required"], nodeData["input"]["optional"])
-				}
-
-				for (const inputName in inputs) {
-					const inputData = inputs[inputName];
-					const type = inputData[0];
-
-					if (!Array.isArray(type)) {
-						types.push(type);
-					}
-				}
-
-				for (const o in nodeData["output"]) {
-					const output = nodeData["output"][o];
-					types.push(output);
-				}
-			}
-
-			return types;
-		}
-
-		function completeColorPalette(colorPalette) {
-			var types = getSlotTypes();
-
-			for (const type of types) {
-				if (!colorPalette.colors.node_slot[type]) {
-					colorPalette.colors.node_slot[type] = "";
-				}
-			}
-
-			colorPalette.colors.node_slot = sortObjectKeys(colorPalette.colors.node_slot);
-
-			return colorPalette;
-		}
-
-		const getColorPaletteTemplate = async () => {
-			let colorPalette = {
-				"id": "my_color_palette_unique_id",
-				"name": "My Color Palette",
-				"colors": {
-					"node_slot": {},
-					"litegraph_base": {},
-					"comfy_base": {}
-				}
-			};
-
-			// Copy over missing keys from default color palette
-			const defaultColorPalette = colorPalettes[defaultColorPaletteId];
-			for (const key in defaultColorPalette.colors.litegraph_base) {
-				if (!colorPalette.colors.litegraph_base[key]) {
-					colorPalette.colors.litegraph_base[key] = "";
-				}
-			}
-			for (const key in defaultColorPalette.colors.comfy_base) {
-				if (!colorPalette.colors.comfy_base[key]) {
-					colorPalette.colors.comfy_base[key] = "";
-				}
-			}
-
-			return completeColorPalette(colorPalette);
-		};
-
-		const getCustomColorPalettes = () => {
-			return app.ui.settings.getSettingValue(idCustomColorPalettes, {});
-		};
-
-		const setCustomColorPalettes = (customColorPalettes) => {
-			return app.ui.settings.setSettingValue(idCustomColorPalettes, customColorPalettes);
-		};
-
-		const addCustomColorPalette = async (colorPalette) => {
-			if (typeof (colorPalette) !== "object") {
-				alert("Invalid color palette.");
-				return;
-			}
-
-			if (!colorPalette.id) {
-				alert("Color palette missing id.");
-				return;
-			}
-
-			if (!colorPalette.name) {
-				alert("Color palette missing name.");
-				return;
-			}
-
-			if (!colorPalette.colors) {
-				alert("Color palette missing colors.");
-				return;
-			}
-
-			if (colorPalette.colors.node_slot && typeof (colorPalette.colors.node_slot) !== "object") {
-				alert("Invalid color palette colors.node_slot.");
-				return;
-			}
-
-			const customColorPalettes = getCustomColorPalettes();
-			customColorPalettes[colorPalette.id] = colorPalette;
-			setCustomColorPalettes(customColorPalettes);
-
-			for (const option of els.select.childNodes) {
-				if (option.value === "custom_" + colorPalette.id) {
-					els.select.removeChild(option);
-				}
-			}
-
-			els.select.append($el("option", {
-				textContent: colorPalette.name + " (custom)",
-				value: "custom_" + colorPalette.id,
-				selected: true
-			}));
-
-			setColorPalette("custom_" + colorPalette.id);
-			await loadColorPalette(colorPalette);
-		};
-
-		const deleteCustomColorPalette = async (colorPaletteId) => {
-			const customColorPalettes = getCustomColorPalettes();
-			delete customColorPalettes[colorPaletteId];
-			setCustomColorPalettes(customColorPalettes);
-
-			for (const option of els.select.childNodes) {
-				if (option.value === defaultColorPaletteId) {
-					option.selected = true;
-				}
-
-				if (option.value === "custom_" + colorPaletteId) {
-					els.select.removeChild(option);
-				}
-			}
-
-			setColorPalette(defaultColorPaletteId);
-			await loadColorPalette(getColorPalette());
-		};
-
-		const loadColorPalette = async (colorPalette) => {
-			colorPalette = await completeColorPalette(colorPalette);
-			if (colorPalette.colors) {
-				// Sets the colors of node slots and links
-				if (colorPalette.colors.node_slot) {
-					Object.assign(app.canvas.default_connection_color_byType, colorPalette.colors.node_slot);
-					Object.assign(LGraphCanvas.link_type_colors, colorPalette.colors.node_slot);
-				}
-				// Sets the colors of the LiteGraph objects
-				if (colorPalette.colors.litegraph_base) {
-					// Everything updates correctly in the loop, except the Node Title and Link Color for some reason
-					app.canvas.node_title_color = colorPalette.colors.litegraph_base.NODE_TITLE_COLOR;
-					app.canvas.default_link_color = colorPalette.colors.litegraph_base.LINK_COLOR;
-
-					for (const key in colorPalette.colors.litegraph_base) {
-						if (colorPalette.colors.litegraph_base.hasOwnProperty(key) && LiteGraph.hasOwnProperty(key)) {
-							LiteGraph[key] = colorPalette.colors.litegraph_base[key];
-						}
-					}
-				}
-				// Sets the color of ComfyUI elements
-				if (colorPalette.colors.comfy_base) {
-					const rootStyle = document.documentElement.style;
-					for (const key in colorPalette.colors.comfy_base) {
-						rootStyle.setProperty('--' + key, colorPalette.colors.comfy_base[key]);
-					}
-				}
-				app.canvas.draw(true, true);
-			}
-		};
-
-		const getColorPalette = (colorPaletteId) => {
-			if (!colorPaletteId) {
-				colorPaletteId = app.ui.settings.getSettingValue(id, defaultColorPaletteId);
-			}
-
-			if (colorPaletteId.startsWith("custom_")) {
-				colorPaletteId = colorPaletteId.substr(7);
-				let customColorPalettes = getCustomColorPalettes();
-				if (customColorPalettes[colorPaletteId]) {
-					return customColorPalettes[colorPaletteId];
-				}
-			}
-
-			return colorPalettes[colorPaletteId];
-		};
-
-		const setColorPalette = (colorPaletteId) => {
-			app.ui.settings.setSettingValue(id, colorPaletteId);
-		};
-
-		const fileInput = $el("input", {
-			type: "file",
-			accept: ".json",
-			style: {display: "none"},
-			parent: document.body,
-			onchange: () => {
-				const file = fileInput.files[0];
-				if (file.type === "application/json" || file.name.endsWith(".json")) {
-					const reader = new FileReader();
-					reader.onload = async () => {
-						await addCustomColorPalette(JSON.parse(reader.result));
-					};
-					reader.readAsText(file);
-				}
-			},
-		});
-
-		app.ui.settings.addSetting({
-			id,
-			name: "Color Palette",
-			type: (name, setter, value) => {
-				const options = [
-					...Object.values(colorPalettes).map(c=> $el("option", {
-						textContent: c.name,
-						value: c.id,
-						selected: c.id === value
-					})),
-					...Object.values(getCustomColorPalettes()).map(c=>$el("option", {
-						textContent: `${c.name} (custom)`,
-						value: `custom_${c.id}`,
-						selected: `custom_${c.id}` === value
-					}))	,
-				];
-
-				els.select = $el("select", {
-					style: {
-						marginBottom: "0.15rem",
-						width: "100%",
-					},
-					onchange: (e) => {
-						setter(e.target.value);
-					}
-				}, options)
-
-				return $el("tr", [
-					$el("td", [
-						$el("label", {
-							for: id.replaceAll(".", "-"),
-							textContent: "Color palette",
-						}),
-					]),
-					$el("td", [
-						els.select,
-						$el("div", {
-							style: {
-								display: "grid",
-								gap: "4px",
-								gridAutoFlow: "column",
-							},
-						}, [
-							$el("input", {
-								type: "button",
-								value: "Export",
-								onclick: async () => {
-									const colorPaletteId = app.ui.settings.getSettingValue(id, defaultColorPaletteId);
-									const colorPalette = await completeColorPalette(getColorPalette(colorPaletteId));
-									const json = JSON.stringify(colorPalette, null, 2); // convert the data to a JSON string
-									const blob = new Blob([json], {type: "application/json"});
-									const url = URL.createObjectURL(blob);
-									const a = $el("a", {
-										href: url,
-										download: colorPaletteId + ".json",
-										style: {display: "none"},
-										parent: document.body,
-									});
-									a.click();
-									setTimeout(function () {
-										a.remove();
-										window.URL.revokeObjectURL(url);
-									}, 0);
-								},
-							}),
-							$el("input", {
-								type: "button",
-								value: "Import",
-								onclick: () => {
-									fileInput.click();
-								}
-							}),
-							$el("input", {
-								type: "button",
-								value: "Template",
-								onclick: async () => {
-									const colorPalette = await getColorPaletteTemplate();
-									const json = JSON.stringify(colorPalette, null, 2); // convert the data to a JSON string
-									const blob = new Blob([json], {type: "application/json"});
-									const url = URL.createObjectURL(blob);
-									const a = $el("a", {
-										href: url,
-										download: "color_palette.json",
-										style: {display: "none"},
-										parent: document.body,
-									});
-									a.click();
-									setTimeout(function () {
-										a.remove();
-										window.URL.revokeObjectURL(url);
-									}, 0);
-								}
-							}),
-							$el("input", {
-								type: "button",
-								value: "Delete",
-								onclick: async () => {
-									let colorPaletteId = app.ui.settings.getSettingValue(id, defaultColorPaletteId);
-
-									if (colorPalettes[colorPaletteId]) {
-										alert("You cannot delete a built-in color palette.");
-										return;
-									}
-
-									if (colorPaletteId.startsWith("custom_")) {
-										colorPaletteId = colorPaletteId.substr(7);
-									}
-
-									await deleteCustomColorPalette(colorPaletteId);
-								}
-							}),
-						]),
-					]),
-				])
-			},
-			defaultValue: defaultColorPaletteId,
-			async onChange(value) {
-				if (!value) {
-					return;
-				}
-
-				let palette = colorPalettes[value];
-				if (palette) {
-					await loadColorPalette(palette);
-				} else if (value.startsWith("custom_")) {
-					value = value.substr(7);
-					let customColorPalettes = getCustomColorPalettes();
-					if (customColorPalettes[value]) {
-						palette = customColorPalettes[value];
-						await loadColorPalette(customColorPalettes[value]);
-					}
-				}
-
-				let {BACKGROUND_IMAGE, CLEAR_BACKGROUND_COLOR} = palette.colors.litegraph_base;
-				if (BACKGROUND_IMAGE === undefined || CLEAR_BACKGROUND_COLOR === undefined) {
-					const base = colorPalettes["dark"].colors.litegraph_base;
-					BACKGROUND_IMAGE = base.BACKGROUND_IMAGE;
-					CLEAR_BACKGROUND_COLOR = base.CLEAR_BACKGROUND_COLOR;
-				}
-				app.canvas.updateBackground(BACKGROUND_IMAGE, CLEAR_BACKGROUND_COLOR);
-			},
-		});
-	},
-});

web/extensions/core/contextMenuFilter.js

@@ -1,148 +0,0 @@
-import {app} from "../../scripts/app.js";
-
-// Adds filtering to combo context menus
-
-const ext = {
-	name: "Comfy.ContextMenuFilter",
-	init() {
-		const ctxMenu = LiteGraph.ContextMenu;
-
-		LiteGraph.ContextMenu = function (values, options) {
-			const ctx = ctxMenu.call(this, values, options);
-
-			// If we are a dark menu (only used for combo boxes) then add a filter input
-			if (options?.className === "dark" && values?.length > 10) {
-				const filter = document.createElement("input");
-				filter.classList.add("comfy-context-menu-filter");
-				filter.placeholder = "Filter list";
-				this.root.prepend(filter);
-
-				const items = Array.from(this.root.querySelectorAll(".litemenu-entry"));
-				let displayedItems = [...items];
-				let itemCount = displayedItems.length;
-
-				// We must request an animation frame for the current node of the active canvas to update.
-				requestAnimationFrame(() => {
-					const currentNode = LGraphCanvas.active_canvas.current_node;
-					const clickedComboValue = currentNode.widgets
-						?.filter(w => w.type === "combo" && w.options.values.length === values.length)
-						.find(w => w.options.values.every((v, i) => v === values[i]))
-						?.value;
-
-					let selectedIndex = clickedComboValue ? values.findIndex(v => v === clickedComboValue) : 0;
-					if (selectedIndex < 0) {
-						selectedIndex = 0;
-					} 
-					let selectedItem = displayedItems[selectedIndex];
-					updateSelected();
-
-					// Apply highlighting to the selected item
-					function updateSelected() {
-						selectedItem?.style.setProperty("background-color", "");
-						selectedItem?.style.setProperty("color", "");
-						selectedItem = displayedItems[selectedIndex];
-						selectedItem?.style.setProperty("background-color", "#ccc", "important");
-						selectedItem?.style.setProperty("color", "#000", "important");
-					}
-
-					const positionList = () => {
-						const rect = this.root.getBoundingClientRect();
-
-						// If the top is off-screen then shift the element with scaling applied
-						if (rect.top < 0) {
-							const scale = 1 - this.root.getBoundingClientRect().height / this.root.clientHeight;
-							const shift = (this.root.clientHeight * scale) / 2;
-							this.root.style.top = -shift + "px";
-						}
-					}
-
-					// Arrow up/down to select items
-					filter.addEventListener("keydown", (event) => {
-						switch (event.key) {
-							case "ArrowUp":
-								event.preventDefault();
-								if (selectedIndex === 0) {
-									selectedIndex = itemCount - 1;
-								} else {
-									selectedIndex--;
-								}
-								updateSelected();
-								break;
-							case "ArrowRight":
-								event.preventDefault();
-								selectedIndex = itemCount - 1;
-								updateSelected();
-								break;
-							case "ArrowDown":
-								event.preventDefault();
-								if (selectedIndex === itemCount - 1) {
-									selectedIndex = 0;
-								} else {
-									selectedIndex++;
-								}
-								updateSelected();
-								break;
-							case "ArrowLeft":
-								event.preventDefault();
-								selectedIndex = 0;
-								updateSelected();
-								break;
-							case "Enter":
-								selectedItem?.click();
-								break;
-							case "Escape":
-								this.close();
-								break;
-						}
-					});
-
-					filter.addEventListener("input", () => {
-						// Hide all items that don't match our filter
-						const term = filter.value.toLocaleLowerCase();
-						// When filtering, recompute which items are visible for arrow up/down and maintain selection.
-						displayedItems = items.filter(item => {
-							const isVisible = !term || item.textContent.toLocaleLowerCase().includes(term);
-							item.style.display = isVisible ? "block" : "none";
-							return isVisible;
-						});
-
-						selectedIndex = 0;
-						if (displayedItems.includes(selectedItem)) {
-							selectedIndex = displayedItems.findIndex(d => d === selectedItem);
-						}
-						itemCount = displayedItems.length;
-
-						updateSelected();
-
-						// If we have an event then we can try and position the list under the source
-						if (options.event) {
-							let top = options.event.clientY - 10;
-
-							const bodyRect = document.body.getBoundingClientRect();
-							const rootRect = this.root.getBoundingClientRect();
-							if (bodyRect.height && top > bodyRect.height - rootRect.height - 10) {
-								top = Math.max(0, bodyRect.height - rootRect.height - 10);
-							}
-
-							this.root.style.top = top + "px";
-							positionList();
-						}
-					});
-
-					requestAnimationFrame(() => {
-						// Focus the filter box when opening
-						filter.focus();
-
-						positionList();
-					});
-				})
-			}
-
-			return ctx;
-		};
-
-		LiteGraph.ContextMenu.prototype = ctxMenu.prototype;
-	},
-}
-
-app.registerExtension(ext);

web/extensions/core/dynamicPrompts.js

@@ -1,48 +0,0 @@
-import { app } from "../../scripts/app.js";
-
-// Allows for simple dynamic prompt replacement
-// Inputs in the format {a|b} will have a random value of a or b chosen when the prompt is queued.
-
-/*
- * Strips C-style line and block comments from a string
- */
-function stripComments(str) {
-	return str.replace(/\/\*[\s\S]*?\*\/|\/\/.*/g,'');
-}
-
-app.registerExtension({
-	name: "Comfy.DynamicPrompts",
-	nodeCreated(node) {
-		if (node.widgets) {
-			// Locate dynamic prompt text widgets
-			// Include any widgets with dynamicPrompts set to true, and customtext
-			const widgets = node.widgets.filter(
-				(n) => n.dynamicPrompts
-			);
-			for (const widget of widgets) {
-				// Override the serialization of the value to resolve dynamic prompts for all widgets supporting it in this node
-				widget.serializeValue = (workflowNode, widgetIndex) => {
-					let prompt = stripComments(widget.value);
-					while (prompt.replace("\\{", "").includes("{") && prompt.replace("\\}", "").includes("}")) {
-						const startIndex = prompt.replace("\\{", "00").indexOf("{");
-						const endIndex = prompt.replace("\\}", "00").indexOf("}");
-
-						const optionsString = prompt.substring(startIndex + 1, endIndex);
-						const options = optionsString.split("|");
-
-						const randomIndex = Math.floor(Math.random() * options.length);
-						const randomOption = options[randomIndex];
-
-						prompt = prompt.substring(0, startIndex) + randomOption + prompt.substring(endIndex + 1);
-					}
-
-					// Overwrite the value in the serialized workflow pnginfo
-					if (workflowNode?.widgets_values)
-						workflowNode.widgets_values[widgetIndex] = prompt;
-
-					return prompt;
-				};
-			}
-		}
-	},
-});

web/extensions/core/editAttention.js

@@ -1,144 +0,0 @@
-import { app } from "../../scripts/app.js";
-
-// Allows you to edit the attention weight by holding ctrl (or cmd) and using the up/down arrow keys
-
-app.registerExtension({
-    name: "Comfy.EditAttention",
-    init() {
-        const editAttentionDelta = app.ui.settings.addSetting({
-            id: "Comfy.EditAttention.Delta",
-            name: "Ctrl+up/down precision",
-            type: "slider",
-            attrs: {
-                min: 0.01,
-                max: 0.5,
-                step: 0.01,
-            },
-            defaultValue: 0.05,
-        });
-
-        function incrementWeight(weight, delta) {
-            const floatWeight = parseFloat(weight);
-            if (isNaN(floatWeight)) return weight;
-            const newWeight = floatWeight + delta;
-            if (newWeight < 0) return "0";
-            return String(Number(newWeight.toFixed(10)));
-        }
-
-        function findNearestEnclosure(text, cursorPos) {
-            let start = cursorPos, end = cursorPos;
-            let openCount = 0, closeCount = 0;
-
-            // Find opening parenthesis before cursor
-            while (start >= 0) {
-                start--;
-                if (text[start] === "(" && openCount === closeCount) break;
-                if (text[start] === "(") openCount++;
-                if (text[start] === ")") closeCount++;
-            }
-            if (start < 0) return false;
-
-            openCount = 0;
-            closeCount = 0;
-
-            // Find closing parenthesis after cursor
-            while (end < text.length) {
-                if (text[end] === ")" && openCount === closeCount) break;
-                if (text[end] === "(") openCount++;
-                if (text[end] === ")") closeCount++;
-                end++;
-            }
-            if (end === text.length) return false;
-
-            return { start: start + 1, end: end };
-        }
-
-        function addWeightToParentheses(text) {
-            const parenRegex = /^\((.*)\)$/;
-            const parenMatch = text.match(parenRegex);
-
-            const floatRegex = /:([+-]?(\d*\.)?\d+([eE][+-]?\d+)?)/;
-            const floatMatch = text.match(floatRegex);
-
-            if (parenMatch && !floatMatch) {
-                return `(${parenMatch[1]}:1.0)`;
-            } else {
-                return text;
-            }
-        };
-
-        function editAttention(event) {
-            const inputField = event.composedPath()[0];
-            const delta = parseFloat(editAttentionDelta.value);
-
-            if (inputField.tagName !== "TEXTAREA") return;
-            if (!(event.key === "ArrowUp" || event.key === "ArrowDown")) return;
-            if (!event.ctrlKey && !event.metaKey) return;
-
-            event.preventDefault();
-
-            let start = inputField.selectionStart;
-            let end = inputField.selectionEnd;
-            let selectedText = inputField.value.substring(start, end);
-
-            // If there is no selection, attempt to find the nearest enclosure, or select the current word
-            if (!selectedText) {
-                const nearestEnclosure = findNearestEnclosure(inputField.value, start);
-                if (nearestEnclosure) {
-                    start = nearestEnclosure.start;
-                    end = nearestEnclosure.end;
-                    selectedText = inputField.value.substring(start, end);
-                } else {
-                    // Select the current word, find the start and end of the word
-                    const delimiters = " .,\\/!?%^*;:{}=-_`~()\r\n\t";
-                    
-                    while (!delimiters.includes(inputField.value[start - 1]) && start > 0) {
-                        start--;
-                    }
-                    
-                    while (!delimiters.includes(inputField.value[end]) && end < inputField.value.length) {
-                        end++;
-                    }
-
-                    selectedText = inputField.value.substring(start, end);
-                    if (!selectedText) return;
-                }
-            }
-
-            // If the selection ends with a space, remove it
-            if (selectedText[selectedText.length - 1] === " ") {
-                selectedText = selectedText.substring(0, selectedText.length - 1);
-                end -= 1;
-            }
-
-            // If there are parentheses left and right of the selection, select them
-            if (inputField.value[start - 1] === "(" && inputField.value[end] === ")") {
-                start -= 1;
-                end += 1;
-                selectedText = inputField.value.substring(start, end);
-            }
-
-            // If the selection is not enclosed in parentheses, add them
-            if (selectedText[0] !== "(" || selectedText[selectedText.length - 1] !== ")") {
-                selectedText = `(${selectedText})`;
-            }
-
-            // If the selection does not have a weight, add a weight of 1.0
-            selectedText = addWeightToParentheses(selectedText);
-
-            // Increment the weight
-            const weightDelta = event.key === "ArrowUp" ? delta : -delta;
-            const updatedText = selectedText.replace(/\((.*):(\d+(?:\.\d+)?)\)/, (match, text, weight) => {
-                weight = incrementWeight(weight, weightDelta);
-                if (weight == 1) {
-                    return text;
-                } else {
-                    return `(${text}:${weight})`;
-                }
-            });
-
-            inputField.setRangeText(updatedText, start, end, "select");
-        }
-        window.addEventListener("keydown", editAttention);
-    },
-});

web/extensions/core/groupNodeManage.js

@@ -1,422 +1,2 @@
-import { $el, ComfyDialog } from "../../scripts/ui.js";
-import { DraggableList } from "../../scripts/ui/draggableList.js";
-import { addStylesheet } from "../../scripts/utils.js";
-import { GroupNodeConfig, GroupNodeHandler } from "./groupNode.js";
-
-addStylesheet(import.meta.url);
-
-const ORDER = Symbol();
-
-function merge(target, source) {
-	if (typeof target === "object" && typeof source === "object") {
-		for (const key in source) {
-			const sv = source[key];
-			if (typeof sv === "object") {
-				let tv = target[key];
-				if (!tv) tv = target[key] = {};
-				merge(tv, source[key]);
-			} else {
-				target[key] = sv;
-			}
-		}
-	}
-
-	return target;
-}
-
-export class ManageGroupDialog extends ComfyDialog {
-	/** @type { Record<"Inputs" | "Outputs" | "Widgets", {tab: HTMLAnchorElement, page: HTMLElement}> } */
-	tabs = {};
-	/** @type { number | null | undefined } */
-	selectedNodeIndex;
-	/** @type { keyof ManageGroupDialog["tabs"] } */
-	selectedTab = "Inputs";
-	/** @type { string | undefined } */
-	selectedGroup;
-
-	/** @type { Record<string, Record<string, Record<string, { name?: string | undefined, visible?: boolean | undefined }>>> } */
-	modifications = {};
-
-	get selectedNodeInnerIndex() {
-		return +this.nodeItems[this.selectedNodeIndex].dataset.nodeindex;
-	}
-
-	constructor(app) {
-		super();
-		this.app = app;
-		this.element = $el("dialog.comfy-group-manage", {
-			parent: document.body,
-		});
-	}
-
-	changeTab(tab) {
-		this.tabs[this.selectedTab].tab.classList.remove("active");
-		this.tabs[this.selectedTab].page.classList.remove("active");
-		this.tabs[tab].tab.classList.add("active");
-		this.tabs[tab].page.classList.add("active");
-		this.selectedTab = tab;
-	}
-
-	changeNode(index, force) {
-		if (!force && this.selectedNodeIndex === index) return;
-
-		if (this.selectedNodeIndex != null) {
-			this.nodeItems[this.selectedNodeIndex].classList.remove("selected");
-		}
-		this.nodeItems[index].classList.add("selected");
-		this.selectedNodeIndex = index;
-
-		if (!this.buildInputsPage() && this.selectedTab === "Inputs") {
-			this.changeTab("Widgets");
-		}
-		if (!this.buildWidgetsPage() && this.selectedTab === "Widgets") {
-			this.changeTab("Outputs");
-		}
-		if (!this.buildOutputsPage() && this.selectedTab === "Outputs") {
-			this.changeTab("Inputs");
-		}
-
-		this.changeTab(this.selectedTab);
-	}
-
-	getGroupData() {
-		this.groupNodeType = LiteGraph.registered_node_types["workflow/" + this.selectedGroup];
-		this.groupNodeDef = this.groupNodeType.nodeData;
-		this.groupData = GroupNodeHandler.getGroupData(this.groupNodeType);
-	}
-
-	changeGroup(group, reset = true) {
-		this.selectedGroup = group;
-		this.getGroupData();
-
-		const nodes = this.groupData.nodeData.nodes;
-		this.nodeItems = nodes.map((n, i) =>
-			$el(
-				"li.draggable-item",
-				{
-					dataset: {
-						nodeindex: n.index + "",
-					},
-					onclick: () => {
-						this.changeNode(i);
-					},
-				},
-				[
-					$el("span.drag-handle"),
-					$el(
-						"div",
-						{
-							textContent: n.title ?? n.type,
-						},
-						n.title
-							? $el("span", {
-									textContent: n.type,
-							  })
-							: []
-					),
-				]
-			)
-		);
-
-		this.innerNodesList.replaceChildren(...this.nodeItems);
-
-		if (reset) {
-			this.selectedNodeIndex = null;
-			this.changeNode(0);
-		} else {
-			const items = this.draggable.getAllItems();
-			let index = items.findIndex(item => item.classList.contains("selected"));
-			if(index === -1) index = this.selectedNodeIndex;
-			this.changeNode(index, true);
-		}
-
-		const ordered = [...nodes];
-		this.draggable?.dispose();
-		this.draggable = new DraggableList(this.innerNodesList, "li");
-		this.draggable.addEventListener("dragend", ({ detail: { oldPosition, newPosition } }) => {
-			if (oldPosition === newPosition) return;
-			ordered.splice(newPosition, 0, ordered.splice(oldPosition, 1)[0]);
-			for (let i = 0; i < ordered.length; i++) {
-				this.storeModification({ nodeIndex: ordered[i].index, section: ORDER, prop: "order", value: i });
-			}
-		});
-	}
-
-	storeModification({ nodeIndex, section, prop, value }) {
-		const groupMod = (this.modifications[this.selectedGroup] ??= {});
-		const nodesMod = (groupMod.nodes ??= {});
-		const nodeMod = (nodesMod[nodeIndex ?? this.selectedNodeInnerIndex] ??= {});
-		const typeMod = (nodeMod[section] ??= {});
-		if (typeof value === "object") {
-			const objMod = (typeMod[prop] ??= {});
-			Object.assign(objMod, value);
-		} else {
-			typeMod[prop] = value;
-		}
-	}
-
-	getEditElement(section, prop, value, placeholder, checked, checkable = true) {
-		if (value === placeholder) value = "";
-
-		const mods = this.modifications[this.selectedGroup]?.nodes?.[this.selectedNodeInnerIndex]?.[section]?.[prop];
-		if (mods) {
-			if (mods.name != null) {
-				value = mods.name;
-			}
-			if (mods.visible != null) {
-				checked = mods.visible;
-			}
-		}
-
-		return $el("div", [
-			$el("input", {
-				value,
-				placeholder,
-				type: "text",
-				onchange: (e) => {
-					this.storeModification({ section, prop, value: { name: e.target.value } });
-				},
-			}),
-			$el("label", { textContent: "Visible" }, [
-				$el("input", {
-					type: "checkbox",
-					checked,
-					disabled: !checkable,
-					onchange: (e) => {
-						this.storeModification({ section, prop, value: { visible: !!e.target.checked } });
-					},
-				}),
-			]),
-		]);
-	}
-
-	buildWidgetsPage() {
-		const widgets = this.groupData.oldToNewWidgetMap[this.selectedNodeInnerIndex];
-		const items = Object.keys(widgets ?? {});
-		const type = app.graph.extra.groupNodes[this.selectedGroup];
-		const config = type.config?.[this.selectedNodeInnerIndex]?.input;
-		this.widgetsPage.replaceChildren(
-			...items.map((oldName) => {
-				return this.getEditElement("input", oldName, widgets[oldName], oldName, config?.[oldName]?.visible !== false);
-			})
-		);
-		return !!items.length;
-	}
-
-	buildInputsPage() {
-		const inputs = this.groupData.nodeInputs[this.selectedNodeInnerIndex];
-		const items = Object.keys(inputs ?? {});
-		const type = app.graph.extra.groupNodes[this.selectedGroup];
-		const config = type.config?.[this.selectedNodeInnerIndex]?.input;
-		this.inputsPage.replaceChildren(
-			...items
-				.map((oldName) => {
-					let value = inputs[oldName];
-					if (!value) {
-						return;
-					}
-
-					return this.getEditElement("input", oldName, value, oldName, config?.[oldName]?.visible !== false);
-				})
-				.filter(Boolean)
-		);
-		return !!items.length;
-	}
-
-	buildOutputsPage() {
-		const nodes = this.groupData.nodeData.nodes;
-		const innerNodeDef = this.groupData.getNodeDef(nodes[this.selectedNodeInnerIndex]);
-		const outputs = innerNodeDef?.output ?? [];
-		const groupOutputs = this.groupData.oldToNewOutputMap[this.selectedNodeInnerIndex];
-
-		const type = app.graph.extra.groupNodes[this.selectedGroup];
-		const config = type.config?.[this.selectedNodeInnerIndex]?.output;
-		const node = this.groupData.nodeData.nodes[this.selectedNodeInnerIndex];
-		const checkable = node.type !== "PrimitiveNode";
-		this.outputsPage.replaceChildren(
-			...outputs
-				.map((type, slot) => {
-					const groupOutputIndex = groupOutputs?.[slot];
-					const oldName = innerNodeDef.output_name?.[slot] ?? type;
-					let value = config?.[slot]?.name;
-					const visible = config?.[slot]?.visible || groupOutputIndex != null;
-					if (!value || value === oldName) {
-						value = "";
-					}
-					return this.getEditElement("output", slot, value, oldName, visible, checkable);
-				})
-				.filter(Boolean)
-		);
-		return !!outputs.length;
-	}
-
-	show(type) {
-		const groupNodes = Object.keys(app.graph.extra?.groupNodes ?? {}).sort((a, b) => a.localeCompare(b));
-
-		this.innerNodesList = $el("ul.comfy-group-manage-list-items");
-		this.widgetsPage = $el("section.comfy-group-manage-node-page");
-		this.inputsPage = $el("section.comfy-group-manage-node-page");
-		this.outputsPage = $el("section.comfy-group-manage-node-page");
-		const pages = $el("div", [this.widgetsPage, this.inputsPage, this.outputsPage]);
-
-		this.tabs = [
-			["Inputs", this.inputsPage],
-			["Widgets", this.widgetsPage],
-			["Outputs", this.outputsPage],
-		].reduce((p, [name, page]) => {
-			p[name] = {
-				tab: $el("a", {
-					onclick: () => {
-						this.changeTab(name);
-					},
-					textContent: name,
-				}),
-				page,
-			};
-			return p;
-		}, {});
-
-		const outer = $el("div.comfy-group-manage-outer", [
-			$el("header", [
-				$el("h2", "Group Nodes"),
-				$el(
-					"select",
-					{
-						onchange: (e) => {
-							this.changeGroup(e.target.value);
-						},
-					},
-					groupNodes.map((g) =>
-						$el("option", {
-							textContent: g,
-							selected: "workflow/" + g === type,
-							value: g,
-						})
-					)
-				),
-			]),
-			$el("main", [
-				$el("section.comfy-group-manage-list", this.innerNodesList),
-				$el("section.comfy-group-manage-node", [
-					$el(
-						"header",
-						Object.values(this.tabs).map((t) => t.tab)
-					),
-					pages,
-				]),
-			]),
-			$el("footer", [
-				$el(
-					"button.comfy-btn",
-					{
-						onclick: (e) => {
-							const node = app.graph._nodes.find((n) => n.type === "workflow/" + this.selectedGroup);
-							if (node) {
-								alert("This group node is in use in the current workflow, please first remove these.");
-								return;
-							}
-							if (confirm(`Are you sure you want to remove the node: "${this.selectedGroup}"`)) {
-								delete app.graph.extra.groupNodes[this.selectedGroup];
-								LiteGraph.unregisterNodeType("workflow/" + this.selectedGroup);
-							}
-							this.show();
-						},
-					},
-					"Delete Group Node"
-				),
-				$el(
-					"button.comfy-btn",
-					{
-						onclick: async () => {
-							let nodesByType;
-							let recreateNodes = [];
-							const types = {};
-							for (const g in this.modifications) {
-								const type = app.graph.extra.groupNodes[g];
-								let config = (type.config ??= {});
-
-								let nodeMods = this.modifications[g]?.nodes;
-								if (nodeMods) {
-									const keys = Object.keys(nodeMods);
-									if (nodeMods[keys[0]][ORDER]) {
-										// If any node is reordered, they will all need sequencing
-										const orderedNodes = [];
-										const orderedMods = {};
-										const orderedConfig = {};
-
-										for (const n of keys) {
-											const order = nodeMods[n][ORDER].order;
-											orderedNodes[order] = type.nodes[+n];
-											orderedMods[order] = nodeMods[n];
-											orderedNodes[order].index = order;
-										}
-
-										// Rewrite links
-										for (const l of type.links) {
-											if (l[0] != null) l[0] = type.nodes[l[0]].index;
-											if (l[2] != null) l[2] = type.nodes[l[2]].index;
-										}
-
-										// Rewrite externals
-										if (type.external) {
-											for (const ext of type.external) {
-												ext[0] = type.nodes[ext[0]];
-											}
-										}
-
-										// Rewrite modifications
-										for (const id of keys) {
-											if (config[id]) {
-												orderedConfig[type.nodes[id].index] = config[id];
-											}
-											delete config[id];
-										}
-
-										type.nodes = orderedNodes;
-										nodeMods = orderedMods;
-										type.config = config = orderedConfig;
-									}
-
-									merge(config, nodeMods);
-								}
-
-								types[g] = type;
-
-								if (!nodesByType) {
-									nodesByType = app.graph._nodes.reduce((p, n) => {
-										p[n.type] ??= [];
-										p[n.type].push(n);
-										return p;
-									}, {});
-								}
-
-								const nodes = nodesByType["workflow/" + g];
-								if (nodes) recreateNodes.push(...nodes);
-							}
-
-							await GroupNodeConfig.registerFromWorkflow(types, {});
-
-							for (const node of recreateNodes) {
-								node.recreate();
-							}
-
-							this.modifications = {};
-							this.app.graph.setDirtyCanvas(true, true);
-							this.changeGroup(this.selectedGroup, false);
-						},
-					},
-					"Save"
-				),
-				$el("button.comfy-btn", { onclick: () => this.element.close() }, "Close"),
-			]),
-		]);
-
-		this.element.replaceChildren(outer);
-		this.changeGroup(type ? groupNodes.find((g) => "workflow/" + g === type) : groupNodes[0]);
-		this.element.showModal();
-
-		this.element.addEventListener("close", () => {
-			this.draggable?.dispose();
-		});
-	}
-}
+// Shim for extensions\core\groupNodeManage.ts
+export const ManageGroupDialog = window.comfyAPI.groupNodeManage.ManageGroupDialog;

web/extensions/core/groupOptions.js

@@ -1,259 +0,0 @@
-import {app} from "../../scripts/app.js";
-
-function setNodeMode(node, mode) {
-    node.mode = mode;
-    node.graph.change();
-}
-
-function addNodesToGroup(group, nodes=[]) {
-    var x1, y1, x2, y2;
-    var nx1, ny1, nx2, ny2;
-    var node;
-
-    x1 = y1 = x2 = y2 = -1;
-    nx1 = ny1 = nx2 = ny2 = -1;
-
-    for (var n of [group._nodes, nodes]) {
-        for (var i in n) {
-            node = n[i]
-
-            nx1 = node.pos[0]
-            ny1 = node.pos[1]
-            nx2 = node.pos[0] + node.size[0]
-            ny2 = node.pos[1] + node.size[1]
-
-            if (node.type != "Reroute") {
-                ny1 -= LiteGraph.NODE_TITLE_HEIGHT;
-            }
-
-            if (node.flags?.collapsed) {
-                ny2 = ny1 + LiteGraph.NODE_TITLE_HEIGHT;
-
-                if (node?._collapsed_width) {
-                    nx2 = nx1 + Math.round(node._collapsed_width);
-                }
-            }
-
-            if (x1 == -1 || nx1 < x1) {
-                x1 = nx1;
-            }
-
-            if (y1 == -1 || ny1 < y1) {
-                y1 = ny1;
-            }
-
-            if (x2 == -1 || nx2 > x2) {
-                x2 = nx2;
-            }
-
-            if (y2 == -1 || ny2 > y2) {
-                y2 = ny2;
-            }
-        }
-    }
-
-    var padding = 10;
-
-    y1 = y1 - Math.round(group.font_size * 1.4);
-
-    group.pos = [x1 - padding, y1 - padding];
-    group.size = [x2 - x1 + padding * 2, y2 - y1 + padding * 2];
-}
-
-app.registerExtension({
-    name: "Comfy.GroupOptions",
-    setup() {
-        const orig = LGraphCanvas.prototype.getCanvasMenuOptions;
-        // graph_mouse
-        LGraphCanvas.prototype.getCanvasMenuOptions = function () {
-            const options = orig.apply(this, arguments);
-            const group = this.graph.getGroupOnPos(this.graph_mouse[0], this.graph_mouse[1]);
-            if (!group) {
-                options.push({
-                    content: "Add Group For Selected Nodes",
-                    disabled: !Object.keys(app.canvas.selected_nodes || {}).length,
-                    callback: () => {
-                        var group = new LiteGraph.LGraphGroup();
-                        addNodesToGroup(group, this.selected_nodes)
-                        app.canvas.graph.add(group);
-                        this.graph.change();
-                    }
-                });
-
-                return options;
-            }
-
-            // Group nodes aren't recomputed until the group is moved, this ensures the nodes are up-to-date
-            group.recomputeInsideNodes();
-            const nodesInGroup = group._nodes;
-
-            options.push({
-                content: "Add Selected Nodes To Group",
-                disabled: !Object.keys(app.canvas.selected_nodes || {}).length,
-                callback: () => {
-                    addNodesToGroup(group, this.selected_nodes)
-                    this.graph.change();
-                }
-            });
-
-            // No nodes in group, return default options
-            if (nodesInGroup.length === 0) {
-                return options;
-            } else {
-                // Add a separator between the default options and the group options
-                options.push(null);
-            }
-
-            // Check if all nodes are the same mode
-            let allNodesAreSameMode = true;
-            for (let i = 1; i < nodesInGroup.length; i++) {
-                if (nodesInGroup[i].mode !== nodesInGroup[0].mode) {
-                    allNodesAreSameMode = false;
-                    break;
-                }
-            }
-
-            options.push({
-                content: "Fit Group To Nodes",
-                callback: () => {
-                    addNodesToGroup(group)
-                    this.graph.change();
-                }
-            });
-
-            options.push({
-                content: "Select Nodes",
-                callback: () => {
-                    this.selectNodes(nodesInGroup);
-                    this.graph.change();
-                    this.canvas.focus();
-                }
-            });
-
-            // Modes
-            // 0: Always
-            // 1: On Event
-            // 2: Never
-            // 3: On Trigger
-            // 4: Bypass
-            // If all nodes are the same mode, add a menu option to change the mode
-            if (allNodesAreSameMode) {
-                const mode = nodesInGroup[0].mode;
-                switch (mode) {
-                    case 0:
-                        // All nodes are always, option to disable, and bypass
-                        options.push({
-                            content: "Set Group Nodes to Never",
-                            callback: () => {
-                                for (const node of nodesInGroup) {
-                                    setNodeMode(node, 2);
-                                }
-                            }
-                        });
-                        options.push({
-                            content: "Bypass Group Nodes",
-                            callback: () => {
-                                for (const node of nodesInGroup) {
-                                    setNodeMode(node, 4);
-                                }
-                            }
-                        });
-                        break;
-                    case 2:
-                        // All nodes are never, option to enable, and bypass
-                        options.push({
-                            content: "Set Group Nodes to Always",
-                            callback: () => {
-                                for (const node of nodesInGroup) {
-                                    setNodeMode(node, 0);
-                                }
-                            }
-                        });
-                        options.push({
-                            content: "Bypass Group Nodes",
-                            callback: () => {
-                                for (const node of nodesInGroup) {
-                                    setNodeMode(node, 4);
-                                }
-                            }
-                        });
-                        break;
-                    case 4:
-                        // All nodes are bypass, option to enable, and disable
-                        options.push({
-                            content: "Set Group Nodes to Always",
-                            callback: () => {
-                                for (const node of nodesInGroup) {
-                                    setNodeMode(node, 0);
-                                }
-                            }
-                        });
-                        options.push({
-                            content: "Set Group Nodes to Never",
-                            callback: () => {
-                                for (const node of nodesInGroup) {
-                                    setNodeMode(node, 2);
-                                }
-                            }
-                        });
-                        break;
-                    default:
-                        // All nodes are On Trigger or On Event(Or other?), option to disable, set to always, or bypass
-                        options.push({
-                            content: "Set Group Nodes to Always",
-                            callback: () => {
-                                for (const node of nodesInGroup) {
-                                    setNodeMode(node, 0);
-                                }
-                            }
-                        });
-                        options.push({
-                            content: "Set Group Nodes to Never",
-                            callback: () => {
-                                for (const node of nodesInGroup) {
-                                    setNodeMode(node, 2);
-                                }
-                            }
-                        });
-                        options.push({
-                            content: "Bypass Group Nodes",
-                            callback: () => {
-                                for (const node of nodesInGroup) {
-                                    setNodeMode(node, 4);
-                                }
-                            }
-                        });
-                        break;
-                }
-            } else {
-                // Nodes are not all the same mode, add a menu option to change the mode to always, never, or bypass
-                options.push({
-                    content: "Set Group Nodes to Always",
-                    callback: () => {
-                        for (const node of nodesInGroup) {
-                            setNodeMode(node, 0);
-                        }
-                    }
-                });
-                options.push({
-                    content: "Set Group Nodes to Never",
-                    callback: () => {
-                        for (const node of nodesInGroup) {
-                            setNodeMode(node, 2);
-                        }
-                    }
-                });
-                options.push({
-                    content: "Bypass Group Nodes",
-                    callback: () => {
-                        for (const node of nodesInGroup) {
-                            setNodeMode(node, 4);
-                        }
-                    }
-                });
-            }
-
-            return options
-        }
-    }
-});

web/extensions/core/invertMenuScrolling.js

@@ -1,36 +0,0 @@
-import { app } from "../../scripts/app.js";
-
-// Inverts the scrolling of context menus
-
-const id = "Comfy.InvertMenuScrolling";
-app.registerExtension({
-	name: id,
-	init() {
-		const ctxMenu = LiteGraph.ContextMenu;
-		const replace = () => {
-			LiteGraph.ContextMenu = function (values, options) {
-				options = options || {};
-				if (options.scroll_speed) {
-					options.scroll_speed *= -1;
-				} else {
-					options.scroll_speed = -0.1;
-				}
-				return ctxMenu.call(this, values, options);
-			};
-			LiteGraph.ContextMenu.prototype = ctxMenu.prototype;
-		};
-		app.ui.settings.addSetting({
-			id,
-			name: "Invert Menu Scrolling",
-			type: "boolean",
-			defaultValue: false,
-			onChange(value) {
-				if (value) {
-					replace();
-				} else {
-					LiteGraph.ContextMenu = ctxMenu;
-				}
-			},
-		});
-	},
-});

web/extensions/core/keybinds.js

@@ -1,69 +0,0 @@
-import {app} from "../../scripts/app.js";
-
-app.registerExtension({
-	name: "Comfy.Keybinds",
-	init() {
-		const keybindListener = function (event) {
-			const modifierPressed = event.ctrlKey || event.metaKey;
-
-			// Queue prompt using ctrl or command + enter
-			if (modifierPressed && event.key === "Enter") {
-				app.queuePrompt(event.shiftKey ? -1 : 0).then();
-				return;
-			}
-
-			const target = event.composedPath()[0];
-			if (["INPUT", "TEXTAREA"].includes(target.tagName)) {
-				return;
-			}
-
-			const modifierKeyIdMap = {
-				s: "#comfy-save-button",
-				o: "#comfy-file-input",
-				Backspace: "#comfy-clear-button",
-				d: "#comfy-load-default-button",
-			};
-
-			const modifierKeybindId = modifierKeyIdMap[event.key];
-			if (modifierPressed && modifierKeybindId) {
-				event.preventDefault();
-
-				const elem = document.querySelector(modifierKeybindId);
-				elem.click();
-				return;
-			}
-
-			// Finished Handling all modifier keybinds, now handle the rest
-			if (event.ctrlKey || event.altKey || event.metaKey) {
-				return;
-			}
-
-			// Close out of modals using escape
-			if (event.key === "Escape") {
-				const modals = document.querySelectorAll(".comfy-modal");
-				const modal = Array.from(modals).find(modal => window.getComputedStyle(modal).getPropertyValue("display") !== "none");
-				if (modal) {
-					modal.style.display = "none";
-				}
-
-				[...document.querySelectorAll("dialog")].forEach(d => {
-					d.close();
-				});
-			}
-
-			const keyIdMap = {
-				q: "#comfy-view-queue-button",
-				h: "#comfy-view-history-button",
-				r: "#comfy-refresh-button",
-			};
-
-			const buttonId = keyIdMap[event.key];
-			if (buttonId) {
-				const button = document.querySelector(buttonId);
-				button.click();
-			}
-		}
-
-		window.addEventListener("keydown", keybindListener, true);
-	}
-});

web/extensions/core/linkRenderMode.js

@@ -1,25 +0,0 @@
-import { app } from "../../scripts/app.js";
-
-const id = "Comfy.LinkRenderMode";
-const ext = {
-	name: id,
-	async setup(app) {
-		app.ui.settings.addSetting({
-			id,
-			name: "Link Render Mode",
-			defaultValue: 2,
-			type: "combo",
-			options: [...LiteGraph.LINK_RENDER_MODES, "Hidden"].map((m, i) => ({
-				value: i,
-				text: m,
-				selected: i == app.canvas.links_render_mode,
-			})),
-			onChange(value) {
-				app.canvas.links_render_mode = +value;
-				app.graph.setDirtyCanvas(true);
-			},
-		});
-	},
-};
-
-app.registerExtension(ext);

web/extensions/core/maskeditor.js

@@ -1,967 +0,0 @@
-import { app } from "../../scripts/app.js";
-import { ComfyDialog, $el } from "../../scripts/ui.js";
-import { ComfyApp } from "../../scripts/app.js";
-import { api } from "../../scripts/api.js"
-import { ClipspaceDialog } from "./clipspace.js";
-
-// Helper function to convert a data URL to a Blob object
-function dataURLToBlob(dataURL) {
-	const parts = dataURL.split(';base64,');
-	const contentType = parts[0].split(':')[1];
-	const byteString = atob(parts[1]);
-	const arrayBuffer = new ArrayBuffer(byteString.length);
-	const uint8Array = new Uint8Array(arrayBuffer);
-	for (let i = 0; i < byteString.length; i++) {
-		uint8Array[i] = byteString.charCodeAt(i);
-	}
-	return new Blob([arrayBuffer], { type: contentType });
-}
-
-function loadedImageToBlob(image) {
-	const canvas = document.createElement('canvas');
-
-	canvas.width = image.width;
-	canvas.height = image.height;
-
-	const ctx = canvas.getContext('2d');
-
-	ctx.drawImage(image, 0, 0);
-
-	const dataURL = canvas.toDataURL('image/png', 1);
-	const blob = dataURLToBlob(dataURL);
-
-	return blob;
-}
-
-function loadImage(imagePath) {
-	return new Promise((resolve, reject) => {
-		const image = new Image();
-
-		image.onload = function() {
-			resolve(image);
-		};
-
-		image.src = imagePath;
-	});
-}
-
-async function uploadMask(filepath, formData) {
-	await api.fetchApi('/upload/mask', {
-		method: 'POST',
-		body: formData
-	}).then(response => {}).catch(error => {
-		console.error('Error:', error);
-	});
-
-	ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']] = new Image();
-	ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src = api.apiURL("/view?" + new URLSearchParams(filepath).toString() + app.getPreviewFormatParam() + app.getRandParam());
-
-	if(ComfyApp.clipspace.images)
-		ComfyApp.clipspace.images[ComfyApp.clipspace['selectedIndex']] = filepath;
-
-	ClipspaceDialog.invalidatePreview();
-}
-
-function prepare_mask(image, maskCanvas, maskCtx, maskColor) {
-	// paste mask data into alpha channel
-	maskCtx.drawImage(image, 0, 0, maskCanvas.width, maskCanvas.height);
-	const maskData = maskCtx.getImageData(0, 0, maskCanvas.width, maskCanvas.height);
-
-	// invert mask
-	for (let i = 0; i < maskData.data.length; i += 4) {
-		if(maskData.data[i+3] == 255)
-			maskData.data[i+3] = 0;
-		else
-			maskData.data[i+3] = 255;
-
-		maskData.data[i] = maskColor.r;
-		maskData.data[i+1] = maskColor.g;
-		maskData.data[i+2] = maskColor.b;
-	}
-
-	maskCtx.globalCompositeOperation = 'source-over';
-	maskCtx.putImageData(maskData, 0, 0);
-}
-
-class MaskEditorDialog extends ComfyDialog {
-	static instance = null;
-
-	static getInstance() {
-		if(!MaskEditorDialog.instance) {
-			MaskEditorDialog.instance = new MaskEditorDialog(app);
-		}
-
-		return MaskEditorDialog.instance;
-	}
-
-	is_layout_created =  false;
-
-	constructor() {
-		super();
-		this.element = $el("div.comfy-modal", { parent: document.body }, 
-			[ $el("div.comfy-modal-content", 
-				[...this.createButtons()]),
-			]);
-	}
-
-	createButtons() {
-		return [];
-	}
-
-	createButton(name, callback) {
-		var button = document.createElement("button");
-		button.style.pointerEvents = "auto";
-		button.innerText = name;
-		button.addEventListener("click", callback);
-		return button;
-	}
-
-	createLeftButton(name, callback) {
-		var button = this.createButton(name, callback);
-		button.style.cssFloat = "left";
-		button.style.marginRight = "4px";
-		return button;
-	}
-
-	createRightButton(name, callback) {
-		var button = this.createButton(name, callback);
-		button.style.cssFloat = "right";
-		button.style.marginLeft = "4px";
-		return button;
-	}
-
-	createLeftSlider(self, name, callback) {
-		const divElement = document.createElement('div');
-		divElement.id = "maskeditor-slider";
-		divElement.style.cssFloat = "left";
-		divElement.style.fontFamily = "sans-serif";
-		divElement.style.marginRight = "4px";
-		divElement.style.color = "var(--input-text)";
-		divElement.style.backgroundColor = "var(--comfy-input-bg)";
-		divElement.style.borderRadius = "8px";
-		divElement.style.borderColor = "var(--border-color)";
-		divElement.style.borderStyle = "solid";
-		divElement.style.fontSize = "15px";
-		divElement.style.height = "21px";
-		divElement.style.padding = "1px 6px";
-		divElement.style.display = "flex";
-		divElement.style.position = "relative";
-		divElement.style.top = "2px";
-		divElement.style.pointerEvents = "auto";
-		self.brush_slider_input = document.createElement('input');
-		self.brush_slider_input.setAttribute('type', 'range');
-		self.brush_slider_input.setAttribute('min', '1');
-		self.brush_slider_input.setAttribute('max', '100');
-		self.brush_slider_input.setAttribute('value', '10');
-		const labelElement = document.createElement("label");
-		labelElement.textContent = name;
-
-		divElement.appendChild(labelElement);
-		divElement.appendChild(self.brush_slider_input);
-
-		self.brush_slider_input.addEventListener("change", callback);
-
-		return divElement;
-	}
-
-	createOpacitySlider(self, name, callback) {
-		const divElement = document.createElement('div');
-		divElement.id = "maskeditor-opacity-slider";
-		divElement.style.cssFloat = "left";
-		divElement.style.fontFamily = "sans-serif";
-		divElement.style.marginRight = "4px";
-		divElement.style.color = "var(--input-text)";
-		divElement.style.backgroundColor = "var(--comfy-input-bg)";
-		divElement.style.borderRadius = "8px";
-		divElement.style.borderColor = "var(--border-color)";
-		divElement.style.borderStyle = "solid";
-		divElement.style.fontSize = "15px";
-		divElement.style.height = "21px";
-		divElement.style.padding = "1px 6px";
-		divElement.style.display = "flex";
-		divElement.style.position = "relative";
-		divElement.style.top = "2px";
-		divElement.style.pointerEvents = "auto";
-		self.opacity_slider_input = document.createElement('input');
-		self.opacity_slider_input.setAttribute('type', 'range');
-		self.opacity_slider_input.setAttribute('min', '0.1');
-		self.opacity_slider_input.setAttribute('max', '1.0');
-		self.opacity_slider_input.setAttribute('step', '0.01')
-		self.opacity_slider_input.setAttribute('value', '0.7');
-		const labelElement = document.createElement("label");
-		labelElement.textContent = name;
-
-		divElement.appendChild(labelElement);
-		divElement.appendChild(self.opacity_slider_input);
-
-		self.opacity_slider_input.addEventListener("input", callback);
-
-		return divElement;
-	}
-
-	setlayout(imgCanvas, maskCanvas) {
-		const self = this;
-
-		// If it is specified as relative, using it only as a hidden placeholder for padding is recommended
-		// to prevent anomalies where it exceeds a certain size and goes outside of the window.
-		var bottom_panel = document.createElement("div");
-		bottom_panel.style.position = "absolute";
-		bottom_panel.style.bottom = "0px";
-		bottom_panel.style.left = "20px";
-		bottom_panel.style.right = "20px";
-		bottom_panel.style.height = "50px";
-		bottom_panel.style.pointerEvents = "none";
-
-		var brush = document.createElement("div");
-		brush.id = "brush";
-		brush.style.backgroundColor = "transparent";
-		brush.style.outline = "1px dashed black";
-		brush.style.boxShadow = "0 0 0 1px white";
-		brush.style.borderRadius = "50%";
-		brush.style.MozBorderRadius = "50%";
-		brush.style.WebkitBorderRadius = "50%";
-		brush.style.position = "absolute";
-		brush.style.zIndex = 8889;
-		brush.style.pointerEvents = "none";
-		this.brush = brush;
-		this.element.appendChild(imgCanvas);
-		this.element.appendChild(maskCanvas);
-		this.element.appendChild(bottom_panel);
-		document.body.appendChild(brush);
-
-		var clearButton = this.createLeftButton("Clear", () => {
-			self.maskCtx.clearRect(0, 0, self.maskCanvas.width, self.maskCanvas.height);
-		});
-		
-		this.brush_size_slider = this.createLeftSlider(self, "Thickness", (event) => {
-			self.brush_size = event.target.value;
-			self.updateBrushPreview(self, null, null);
-		});
-
-		this.brush_opacity_slider = this.createOpacitySlider(self, "Opacity", (event) => {
-			self.brush_opacity = event.target.value;
-			if (self.brush_color_mode !== "negative") {
-			    self.maskCanvas.style.opacity = self.brush_opacity;
-			}
-		});
-
-		this.colorButton = this.createLeftButton(this.getColorButtonText(), () => {
-			if (self.brush_color_mode === "black") {
-				self.brush_color_mode = "white";
-			}
-			else if (self.brush_color_mode === "white") {
-				self.brush_color_mode = "negative";
-			}
-			else {
-				self.brush_color_mode = "black";
-			}
-
-			self.updateWhenBrushColorModeChanged();
-		});
-
-		var cancelButton = this.createRightButton("Cancel", () => {
-			document.removeEventListener("mouseup", MaskEditorDialog.handleMouseUp);
-			document.removeEventListener("keydown", MaskEditorDialog.handleKeyDown);
-			self.close();
-		});
-
-		this.saveButton = this.createRightButton("Save", () => {
-			document.removeEventListener("mouseup", MaskEditorDialog.handleMouseUp);
-			document.removeEventListener("keydown", MaskEditorDialog.handleKeyDown);
-				self.save();
-			});
-
-		this.element.appendChild(imgCanvas);
-		this.element.appendChild(maskCanvas);
-		this.element.appendChild(bottom_panel);
-
-		bottom_panel.appendChild(clearButton);
-		bottom_panel.appendChild(this.saveButton);
-		bottom_panel.appendChild(cancelButton);
-		bottom_panel.appendChild(this.brush_size_slider);
-		bottom_panel.appendChild(this.brush_opacity_slider);
-		bottom_panel.appendChild(this.colorButton);
-
-		imgCanvas.style.position = "absolute";
-		maskCanvas.style.position = "absolute";
-
-		imgCanvas.style.top = "200";
-		imgCanvas.style.left = "0";
-
-		maskCanvas.style.top = imgCanvas.style.top;
-		maskCanvas.style.left = imgCanvas.style.left;
-
-		const maskCanvasStyle = this.getMaskCanvasStyle();
-		maskCanvas.style.mixBlendMode = maskCanvasStyle.mixBlendMode;
-		maskCanvas.style.opacity = maskCanvasStyle.opacity;
-	}
-
-	async show() {
-		this.zoom_ratio = 1.0;
-		this.pan_x = 0;
-		this.pan_y = 0;
-
-		if(!this.is_layout_created) {
-			// layout
-			const imgCanvas = document.createElement('canvas');
-			const maskCanvas = document.createElement('canvas');
-
-			imgCanvas.id = "imageCanvas";
-			maskCanvas.id = "maskCanvas";
-
-			this.setlayout(imgCanvas, maskCanvas);
-
-			// prepare content
-			this.imgCanvas = imgCanvas;
-			this.maskCanvas = maskCanvas;
-			this.maskCtx = maskCanvas.getContext('2d', {willReadFrequently: true });
-
-			this.setEventHandler(maskCanvas);
-
-			this.is_layout_created = true;
-
-			// replacement of onClose hook since close is not real close
-			const self = this;
-			const observer = new MutationObserver(function(mutations) {
-			mutations.forEach(function(mutation) {
-					if (mutation.type === 'attributes' && mutation.attributeName === 'style') {
-						if(self.last_display_style && self.last_display_style != 'none' && self.element.style.display == 'none') {
-							document.removeEventListener("mouseup", MaskEditorDialog.handleMouseUp);
-							self.brush.style.display = "none";
-							ComfyApp.onClipspaceEditorClosed();
-						}
-
-						self.last_display_style = self.element.style.display;
-					}
-				});
-			});
-
-			const config = { attributes: true };
-			observer.observe(this.element, config);
-		}
-
-		// The keydown event needs to be reconfigured when closing the dialog as it gets removed.
-		document.addEventListener('keydown', MaskEditorDialog.handleKeyDown);
-
-		if(ComfyApp.clipspace_return_node) {
-			this.saveButton.innerText = "Save to node";
-		}
-		else {
-			this.saveButton.innerText = "Save";
-		}
-		this.saveButton.disabled = false;
-
-		this.element.style.display = "block";
-		this.element.style.width = "85%";
-		this.element.style.margin = "0 7.5%";
-		this.element.style.height = "100vh";
-		this.element.style.top = "50%";
-		this.element.style.left = "42%";
-		this.element.style.zIndex = 8888; // NOTE: alert dialog must be high priority.
-
-		await this.setImages(this.imgCanvas);
-
-		this.is_visible = true;
-	}
-
-	isOpened() {
-		return this.element.style.display == "block";
-	}
-
-	invalidateCanvas(orig_image, mask_image) {
-		this.imgCanvas.width = orig_image.width;
-		this.imgCanvas.height = orig_image.height;
-
-		this.maskCanvas.width = orig_image.width;
-		this.maskCanvas.height = orig_image.height;
-
-		let imgCtx = this.imgCanvas.getContext('2d', {willReadFrequently: true });
-		let maskCtx = this.maskCanvas.getContext('2d', {willReadFrequently: true });
-
-		imgCtx.drawImage(orig_image, 0, 0, orig_image.width, orig_image.height);
-		prepare_mask(mask_image, this.maskCanvas, maskCtx, this.getMaskColor());
-	}
-
-	async setImages(imgCanvas) {
-		let self = this;
-
-		const imgCtx = imgCanvas.getContext('2d', {willReadFrequently: true });
-		const maskCtx = this.maskCtx;
-		const maskCanvas = this.maskCanvas;
-
-		imgCtx.clearRect(0,0,this.imgCanvas.width,this.imgCanvas.height);
-		maskCtx.clearRect(0,0,this.maskCanvas.width,this.maskCanvas.height);
-
-		// image load
-		const filepath = ComfyApp.clipspace.images;
-
-		const alpha_url = new URL(ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src)
-		alpha_url.searchParams.delete('channel');
-		alpha_url.searchParams.delete('preview');
-		alpha_url.searchParams.set('channel', 'a');
-		let mask_image = await loadImage(alpha_url);
-
-		// original image load
-		const rgb_url = new URL(ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src);
-		rgb_url.searchParams.delete('channel');
-		rgb_url.searchParams.set('channel', 'rgb');
-		this.image = new Image();
-		this.image.onload = function() {
-			maskCanvas.width = self.image.width;
-			maskCanvas.height = self.image.height;
-
-			self.invalidateCanvas(self.image, mask_image);
-			self.initializeCanvasPanZoom();
-		};
-		this.image.src = rgb_url;
-	}
-
-	initializeCanvasPanZoom() {
-		// set initialize
-		let drawWidth = this.image.width;
-		let drawHeight = this.image.height;
-
-		let width = this.element.clientWidth;
-		let height = this.element.clientHeight;
-
-		if (this.image.width > width) {
-			drawWidth = width;
-			drawHeight = (drawWidth / this.image.width) * this.image.height;
-		}
-
-		if (drawHeight > height) {
-			drawHeight = height;
-			drawWidth = (drawHeight / this.image.height) * this.image.width;
-		}
-
-		this.zoom_ratio = drawWidth/this.image.width;
-
-		const canvasX = (width - drawWidth) / 2;
-		const canvasY = (height - drawHeight) / 2;
-		this.pan_x = canvasX;
-		this.pan_y = canvasY;
-
-		this.invalidatePanZoom();
-	}
-
-
-	invalidatePanZoom() {
-		let raw_width = this.image.width * this.zoom_ratio;
-		let raw_height = this.image.height * this.zoom_ratio;
-
-		if(this.pan_x + raw_width < 10) {
-			this.pan_x = 10 - raw_width;
-		}
-
-		if(this.pan_y + raw_height < 10) {
-			this.pan_y = 10 - raw_height;
-		}
-
-		let width = `${raw_width}px`;
-		let height = `${raw_height}px`;
-
-		let left = `${this.pan_x}px`;
-		let top = `${this.pan_y}px`;
-
-		this.maskCanvas.style.width = width;
-		this.maskCanvas.style.height = height;
-		this.maskCanvas.style.left = left;
-		this.maskCanvas.style.top = top;
-
-		this.imgCanvas.style.width = width;
-		this.imgCanvas.style.height = height;
-		this.imgCanvas.style.left = left;
-		this.imgCanvas.style.top = top;
-	}
-
-
-	setEventHandler(maskCanvas) {
-		const self = this;
-
-		if(!this.handler_registered) {
-			maskCanvas.addEventListener("contextmenu", (event) => {
-				event.preventDefault();
-			});
-
-			this.element.addEventListener('wheel', (event) => this.handleWheelEvent(self,event));
-			this.element.addEventListener('pointermove', (event) => this.pointMoveEvent(self,event));
-			this.element.addEventListener('touchmove', (event) => this.pointMoveEvent(self,event));
-
-			this.element.addEventListener('dragstart', (event) => {
-				if(event.ctrlKey) {
-					event.preventDefault();
-				}
-			});
-
-			maskCanvas.addEventListener('pointerdown', (event) => this.handlePointerDown(self,event));
-			maskCanvas.addEventListener('pointermove', (event) => this.draw_move(self,event));
-			maskCanvas.addEventListener('touchmove', (event) => this.draw_move(self,event));
-			maskCanvas.addEventListener('pointerover', (event) => { this.brush.style.display = "block"; });
-			maskCanvas.addEventListener('pointerleave', (event) => { this.brush.style.display = "none"; });
-
-			document.addEventListener('pointerup', MaskEditorDialog.handlePointerUp);
-
-			this.handler_registered = true;
-		}
-	}
-
-	getMaskCanvasStyle() {
-		if (this.brush_color_mode === "negative") {
-			return {
-				mixBlendMode: "difference",
-				opacity: "1",
-			};
-		}
-		else {
-			return {
-				mixBlendMode: "initial",
-				opacity: this.brush_opacity,
-			};
-		}
-	}
-
-	getMaskColor() {
-		if (this.brush_color_mode === "black") {
-			return { r: 0, g: 0, b: 0 };
-		}
-		if (this.brush_color_mode === "white") {
-			return { r: 255, g: 255, b: 255 };
-		}
-		if (this.brush_color_mode === "negative") {
-			// negative effect only works with white color
-			return { r: 255, g: 255, b: 255 };
-		}
-
-		return { r: 0, g: 0, b: 0 };
-	}
-
-	getMaskFillStyle() {
-		const maskColor = this.getMaskColor();
-
-		return "rgb(" + maskColor.r + "," + maskColor.g + "," + maskColor.b + ")";
-	}
-
-	getColorButtonText() {
-		let colorCaption = "unknown";
-
-		if (this.brush_color_mode === "black") {
-			colorCaption = "black";
-		}
-		else if (this.brush_color_mode === "white") {
-			colorCaption = "white";
-		}
-		else if (this.brush_color_mode === "negative") {
-			colorCaption = "negative";
-		}
-
-		return "Color: " + colorCaption;
-	}
-
-	updateWhenBrushColorModeChanged() {
-		this.colorButton.innerText = this.getColorButtonText();
-
-		// update mask canvas css styles
-
-		const maskCanvasStyle = this.getMaskCanvasStyle();
-		this.maskCanvas.style.mixBlendMode = maskCanvasStyle.mixBlendMode;
-		this.maskCanvas.style.opacity = maskCanvasStyle.opacity;
-
-		// update mask canvas rgb colors
-
-		const maskColor = this.getMaskColor();
-
-		const maskData = this.maskCtx.getImageData(0, 0, this.maskCanvas.width, this.maskCanvas.height);
-		
-		for (let i = 0; i < maskData.data.length; i += 4) {
-			maskData.data[i] = maskColor.r;
-			maskData.data[i+1] = maskColor.g;
-			maskData.data[i+2] = maskColor.b;
-		}
-	
-		this.maskCtx.putImageData(maskData, 0, 0);
-	}
-
-	brush_opacity = 0.7;
-	brush_size = 10;
-	brush_color_mode = "black";
-	drawing_mode = false;
-	lastx = -1;
-	lasty = -1;
-	lasttime = 0;
-
-	static handleKeyDown(event) {
-		const self = MaskEditorDialog.instance;
-		if (event.key === ']') {
-			self.brush_size = Math.min(self.brush_size+2, 100);
-			self.brush_slider_input.value = self.brush_size;
-		} else if (event.key === '[') {
-			self.brush_size = Math.max(self.brush_size-2, 1);
-			self.brush_slider_input.value = self.brush_size;
-		} else if(event.key === 'Enter') {
-			self.save();
-		}
-
-		self.updateBrushPreview(self);
-	}
-
-	static handlePointerUp(event) {
-		event.preventDefault();
-
-		this.mousedown_x = null;
-		this.mousedown_y = null;
-
-		MaskEditorDialog.instance.drawing_mode = false;
-	}
-
-	updateBrushPreview(self) {
-		const brush = self.brush;
-
-		var centerX = self.cursorX;
-		var centerY = self.cursorY;
-
-		brush.style.width = self.brush_size * 2 * this.zoom_ratio + "px";
-		brush.style.height = self.brush_size * 2 * this.zoom_ratio + "px";
-		brush.style.left = (centerX - self.brush_size * this.zoom_ratio) + "px";
-		brush.style.top = (centerY - self.brush_size * this.zoom_ratio) + "px";
-	}
-
-	handleWheelEvent(self, event) {
-		event.preventDefault();
-
-		if(event.ctrlKey) {
-			// zoom canvas
-			if(event.deltaY < 0) {
-				this.zoom_ratio = Math.min(10.0, this.zoom_ratio+0.2);
-			}
-			else {
-				this.zoom_ratio = Math.max(0.2, this.zoom_ratio-0.2);
-			}
-
-			this.invalidatePanZoom();
-		}
-		else {
-			// adjust brush size
-			if(event.deltaY < 0)
-				this.brush_size = Math.min(this.brush_size+2, 100);
-			else
-				this.brush_size = Math.max(this.brush_size-2, 1);
-
-			this.brush_slider_input.value = this.brush_size;
-
-			this.updateBrushPreview(this);
-		}
-	}
-
-	pointMoveEvent(self, event) {
-		this.cursorX = event.pageX;
-		this.cursorY = event.pageY;
-
-		self.updateBrushPreview(self);
-
-		if(event.ctrlKey) {
-			event.preventDefault();
-			self.pan_move(self, event);
-		}
-
-		let left_button_down = window.TouchEvent && event instanceof TouchEvent || event.buttons == 1;
-
-		if(event.shiftKey && left_button_down) {
-			self.drawing_mode = false;
-
-			const y = event.clientY;
-			let delta = (self.zoom_lasty - y)*0.005;
-			self.zoom_ratio = Math.max(Math.min(10.0, self.last_zoom_ratio - delta), 0.2);
-
-			this.invalidatePanZoom();
-			return;
-		}
-	}
-
-	pan_move(self, event) {
-		if(event.buttons == 1) {
-			if(this.mousedown_x) {
-				let deltaX = this.mousedown_x - event.clientX;
-				let deltaY = this.mousedown_y - event.clientY;
-
-				self.pan_x = this.mousedown_pan_x - deltaX;
-				self.pan_y = this.mousedown_pan_y - deltaY;
-
-				self.invalidatePanZoom();
-			}
-		}
-	}
-
-	draw_move(self, event) {
-		if(event.ctrlKey || event.shiftKey) {
-			return;
-		}
-
-		event.preventDefault();
-
-		this.cursorX = event.pageX;
-		this.cursorY = event.pageY;
-
-		self.updateBrushPreview(self);
-
-		let left_button_down = window.TouchEvent && event instanceof TouchEvent || event.buttons == 1;
-		let right_button_down = [2, 5, 32].includes(event.buttons);
-
-		if (!event.altKey && left_button_down) {
-			var diff = performance.now() - self.lasttime;
-
-			const maskRect = self.maskCanvas.getBoundingClientRect();
-
-			var x = event.offsetX;
-			var y = event.offsetY
-
-			if(event.offsetX == null) {
-				x = event.targetTouches[0].clientX - maskRect.left;
-			}
-
-			if(event.offsetY == null) {
-				y = event.targetTouches[0].clientY - maskRect.top;
-			}
-
-			x /= self.zoom_ratio;
-			y /= self.zoom_ratio;
-
-			var brush_size = this.brush_size;
-			if(event instanceof PointerEvent && event.pointerType == 'pen') {
-				brush_size *= event.pressure;
-				this.last_pressure = event.pressure;
-			}
-			else if(window.TouchEvent && event instanceof TouchEvent && diff < 20){
-				// The firing interval of PointerEvents in Pen is unreliable, so it is supplemented by TouchEvents.
-				brush_size *= this.last_pressure;
-			}
-			else {
-				brush_size = this.brush_size;
-			}
-
-			if(diff > 20 && !this.drawing_mode)
-				requestAnimationFrame(() => {
-					self.maskCtx.beginPath();
-					self.maskCtx.fillStyle = this.getMaskFillStyle();
-					self.maskCtx.globalCompositeOperation = "source-over";
-					self.maskCtx.arc(x, y, brush_size, 0, Math.PI * 2, false);
-					self.maskCtx.fill();
-					self.lastx = x;
-					self.lasty = y;
-				});
-			else
-				requestAnimationFrame(() => {
-					self.maskCtx.beginPath();
-					self.maskCtx.fillStyle = this.getMaskFillStyle();
-					self.maskCtx.globalCompositeOperation = "source-over";
-
-					var dx = x - self.lastx;
-					var dy = y - self.lasty;
-
-					var distance = Math.sqrt(dx * dx + dy * dy);
-					var directionX = dx / distance;
-					var directionY = dy / distance;
-
-					for (var i = 0; i < distance; i+=5) {
-						var px = self.lastx + (directionX * i);
-						var py = self.lasty + (directionY * i);
-						self.maskCtx.arc(px, py, brush_size, 0, Math.PI * 2, false);
-						self.maskCtx.fill();
-					}
-					self.lastx = x;
-					self.lasty = y;
-				});
-
-			self.lasttime = performance.now();
-		}
-		else if((event.altKey && left_button_down) || right_button_down) {
-			const maskRect = self.maskCanvas.getBoundingClientRect();
-			const x = (event.offsetX || event.targetTouches[0].clientX - maskRect.left) / self.zoom_ratio;
-			const y = (event.offsetY || event.targetTouches[0].clientY - maskRect.top) / self.zoom_ratio;
-
-			var brush_size = this.brush_size;
-			if(event instanceof PointerEvent && event.pointerType == 'pen') {
-				brush_size *= event.pressure;
-				this.last_pressure = event.pressure;
-			}
-			else if(window.TouchEvent && event instanceof TouchEvent && diff < 20){
-				brush_size *= this.last_pressure;
-			}
-			else {
-				brush_size = this.brush_size;
-			}
-
-			if(diff > 20 && !drawing_mode) // cannot tracking drawing_mode for touch event
-				requestAnimationFrame(() => {
-					self.maskCtx.beginPath();
-					self.maskCtx.globalCompositeOperation = "destination-out";
-					self.maskCtx.arc(x, y, brush_size, 0, Math.PI * 2, false);
-					self.maskCtx.fill();
-					self.lastx = x;
-					self.lasty = y;
-				});
-			else
-				requestAnimationFrame(() => {
-					self.maskCtx.beginPath();
-					self.maskCtx.globalCompositeOperation = "destination-out";
-					
-					var dx = x - self.lastx;
-					var dy = y - self.lasty;
-
-					var distance = Math.sqrt(dx * dx + dy * dy);
-					var directionX = dx / distance;
-					var directionY = dy / distance;
-
-					for (var i = 0; i < distance; i+=5) {
-						var px = self.lastx + (directionX * i);
-						var py = self.lasty + (directionY * i);
-						self.maskCtx.arc(px, py, brush_size, 0, Math.PI * 2, false);
-						self.maskCtx.fill();
-					}
-					self.lastx = x;
-					self.lasty = y;
-				});
-
-				self.lasttime = performance.now();
-		}
-	}
-
-	handlePointerDown(self, event) {
-		if(event.ctrlKey) {
-			if (event.buttons == 1) {
-				this.mousedown_x = event.clientX;
-				this.mousedown_y = event.clientY;
-
-				this.mousedown_pan_x = this.pan_x;
-				this.mousedown_pan_y = this.pan_y;
-			}
-			return;
-		}
-
-		var brush_size = this.brush_size;
-		if(event instanceof PointerEvent && event.pointerType == 'pen') {
-			brush_size *= event.pressure;
-			this.last_pressure = event.pressure;
-		}
-
-		if ([0, 2, 5].includes(event.button)) {
-			self.drawing_mode = true;
-
-			event.preventDefault();
-
-			if(event.shiftKey) {
-				self.zoom_lasty = event.clientY;
-				self.last_zoom_ratio = self.zoom_ratio;
-				return;
-			}
-
-			const maskRect = self.maskCanvas.getBoundingClientRect();
-			const x = (event.offsetX || event.targetTouches[0].clientX - maskRect.left) / self.zoom_ratio;
-			const y = (event.offsetY || event.targetTouches[0].clientY - maskRect.top) / self.zoom_ratio;
-
-			self.maskCtx.beginPath();
-			if (!event.altKey && event.button == 0) {
-				self.maskCtx.fillStyle = this.getMaskFillStyle();
-				self.maskCtx.globalCompositeOperation = "source-over";
-			} else {
-				self.maskCtx.globalCompositeOperation = "destination-out";
-			}
-			self.maskCtx.arc(x, y, brush_size, 0, Math.PI * 2, false);
-			self.maskCtx.fill();
-			self.lastx = x;
-			self.lasty = y;
-			self.lasttime = performance.now();
-		}
-	}
-
-	async save() {
-		const backupCanvas = document.createElement('canvas');
-		const backupCtx = backupCanvas.getContext('2d', {willReadFrequently:true});
-		backupCanvas.width = this.image.width;
-		backupCanvas.height = this.image.height;
-
-		backupCtx.clearRect(0,0, backupCanvas.width, backupCanvas.height);
-		backupCtx.drawImage(this.maskCanvas,
-			0, 0, this.maskCanvas.width, this.maskCanvas.height,
-			0, 0, backupCanvas.width, backupCanvas.height);
-
-		// paste mask data into alpha channel
-		const backupData = backupCtx.getImageData(0, 0, backupCanvas.width, backupCanvas.height);
-
-		// refine mask image
-		for (let i = 0; i < backupData.data.length; i += 4) {
-			if(backupData.data[i+3] == 255)
-				backupData.data[i+3] = 0;
-			else
-				backupData.data[i+3] = 255;
-
-			backupData.data[i] = 0;
-			backupData.data[i+1] = 0;
-			backupData.data[i+2] = 0;
-		}
-
-		backupCtx.globalCompositeOperation = 'source-over';
-		backupCtx.putImageData(backupData, 0, 0);
-
-		const formData = new FormData();
-		const filename = "clipspace-mask-" + performance.now() + ".png";
-
-		const item =
-			{
-				"filename": filename,
-				"subfolder": "clipspace",
-				"type": "input",
-			};
-
-		if(ComfyApp.clipspace.images)
-			ComfyApp.clipspace.images[0] = item;
-
-		if(ComfyApp.clipspace.widgets) {
-			const index = ComfyApp.clipspace.widgets.findIndex(obj => obj.name === 'image');
-
-			if(index >= 0)
-				ComfyApp.clipspace.widgets[index].value = item;
-		}
-
-		const dataURL = backupCanvas.toDataURL();
-		const blob = dataURLToBlob(dataURL);
-
-		let original_url = new URL(this.image.src);
-
-		const original_ref = { filename: original_url.searchParams.get('filename') };
-
-		let original_subfolder = original_url.searchParams.get("subfolder");
-		if(original_subfolder)
-			original_ref.subfolder = original_subfolder;
-
-		let original_type = original_url.searchParams.get("type");
-		if(original_type)
-			original_ref.type = original_type;
-
-		formData.append('image', blob, filename);
-		formData.append('original_ref', JSON.stringify(original_ref));
-		formData.append('type', "input");
-		formData.append('subfolder', "clipspace");
-
-		this.saveButton.innerText = "Saving...";
-		this.saveButton.disabled = true;
-		await uploadMask(item, formData);
-		ComfyApp.onClipspaceEditorSave();
-		this.close();
-	}
-}
-
-app.registerExtension({
-	name: "Comfy.MaskEditor",
-	init(app) {
-		ComfyApp.open_maskeditor =
-			function () {
-				const dlg = MaskEditorDialog.getInstance();
-				if(!dlg.isOpened()) {
-					dlg.show();
-				}
-			};
-
-		const context_predicate = () => ComfyApp.clipspace && ComfyApp.clipspace.imgs && ComfyApp.clipspace.imgs.length > 0
-		ClipspaceDialog.registerButton("MaskEditor", context_predicate, ComfyApp.open_maskeditor);
-	}
-});

web/extensions/core/nodeTemplates.js

@@ -1,412 +0,0 @@
-import { app } from "../../scripts/app.js";
-import { api } from "../../scripts/api.js";
-import { ComfyDialog, $el } from "../../scripts/ui.js";
-import { GroupNodeConfig, GroupNodeHandler } from "./groupNode.js";
-
-// Adds the ability to save and add multiple nodes as a template
-// To save:
-// Select multiple nodes (ctrl + drag to select a region or ctrl+click individual nodes)
-// Right click the canvas
-// Save Node Template -> give it a name
-//
-// To add:
-// Right click the canvas
-// Node templates -> click the one to add
-//
-// To delete/rename:
-// Right click the canvas
-// Node templates -> Manage
-//
-// To rearrange:
-// Open the manage dialog and Drag and drop elements using the "Name:" label as handle
-
-const id = "Comfy.NodeTemplates";
-const file = "comfy.templates.json";
-
-class ManageTemplates extends ComfyDialog {
-	constructor() {
-		super();
-		this.load().then((v) => {
-			this.templates = v;
-		});
-
-		this.element.classList.add("comfy-manage-templates");
-		this.draggedEl = null;
-		this.saveVisualCue = null;
-		this.emptyImg = new Image();
-		this.emptyImg.src = "data:image/gif;base64,R0lGODlhAQABAIAAAAUEBAAAACwAAAAAAQABAAACAkQBADs=";
-
-		this.importInput = $el("input", {
-			type: "file",
-			accept: ".json",
-			multiple: true,
-			style: { display: "none" },
-			parent: document.body,
-			onchange: () => this.importAll(),
-		});
-	}
-
-	createButtons() {
-		const btns = super.createButtons();
-		btns[0].textContent = "Close";
-		btns[0].onclick = (e) => {
-			clearTimeout(this.saveVisualCue);
-			this.close();
-		};
-		btns.unshift(
-			$el("button", {
-				type: "button",
-				textContent: "Export",
-				onclick: () => this.exportAll(),
-			})
-		);
-		btns.unshift(
-			$el("button", {
-				type: "button",
-				textContent: "Import",
-				onclick: () => {
-					this.importInput.click();
-				},
-			})
-		);
-		return btns;
-	}
-
-	async load() {
-		let templates = [];
-		if (app.storageLocation === "server") {
-			if (app.isNewUserSession) {
-				// New user so migrate existing templates
-				const json = localStorage.getItem(id);
-				if (json) {
-					templates = JSON.parse(json);
-				}
-				await api.storeUserData(file, json, { stringify: false });
-			} else {
-				const res = await api.getUserData(file);
-				if (res.status === 200) {
-					try {
-						templates = await res.json();
-					} catch (error) {
-					}
-				} else if (res.status !== 404) {
-					console.error(res.status + " " + res.statusText);
-				}
-			}
-		} else {
-			const json = localStorage.getItem(id);
-			if (json) {
-				templates = JSON.parse(json);
-			}
-		}
-
-		return templates ?? [];
-	}
-
-	async store() {
-		if(app.storageLocation === "server") {
-			const templates = JSON.stringify(this.templates, undefined, 4);
-			localStorage.setItem(id, templates); // Backwards compatibility
-			try {
-				await api.storeUserData(file, templates, { stringify: false });
-			} catch (error) {
-				console.error(error);
-				alert(error.message);
-			}
-		} else {
-			localStorage.setItem(id, JSON.stringify(this.templates));
-		}
-	}
-
-	async importAll() {
-		for (const file of this.importInput.files) {
-			if (file.type === "application/json" || file.name.endsWith(".json")) {
-				const reader = new FileReader();
-				reader.onload = async () => {
-					const importFile = JSON.parse(reader.result);
-					if (importFile?.templates) {
-						for (const template of importFile.templates) {
-							if (template?.name && template?.data) {
-								this.templates.push(template);
-							}
-						}
-						await this.store();
-					}
-				};
-				await reader.readAsText(file);
-			}
-		}
-
-		this.importInput.value = null;
-
-		this.close();
-	}
-
-	exportAll() {
-		if (this.templates.length == 0) {
-			alert("No templates to export.");
-			return;
-		}
-
-		const json = JSON.stringify({ templates: this.templates }, null, 2); // convert the data to a JSON string
-		const blob = new Blob([json], { type: "application/json" });
-		const url = URL.createObjectURL(blob);
-		const a = $el("a", {
-			href: url,
-			download: "node_templates.json",
-			style: { display: "none" },
-			parent: document.body,
-		});
-		a.click();
-		setTimeout(function () {
-			a.remove();
-			window.URL.revokeObjectURL(url);
-		}, 0);
-	}
-
-	show() {
-		// Show list of template names + delete button
-		super.show(
-			$el(
-				"div",
-				{},
-				this.templates.flatMap((t,i) => {
-					let nameInput;
-					return [
-						$el(
-							"div",
-							{
-								dataset: { id: i },
-								className: "tempateManagerRow",
-								style: {
-									display: "grid",
-									gridTemplateColumns: "1fr auto",
-									border: "1px dashed transparent",
-									gap: "5px",
-									backgroundColor: "var(--comfy-menu-bg)"
-								},
-								ondragstart: (e) => {
-									this.draggedEl = e.currentTarget;
-									e.currentTarget.style.opacity = "0.6";
-									e.currentTarget.style.border = "1px dashed yellow";
-									e.dataTransfer.effectAllowed = 'move';
-									e.dataTransfer.setDragImage(this.emptyImg, 0, 0);
-								},
-								ondragend: (e) => {
-									e.target.style.opacity = "1";
-									e.currentTarget.style.border = "1px dashed transparent";
-									e.currentTarget.removeAttribute("draggable");
-
-									// rearrange the elements
-									this.element.querySelectorAll('.tempateManagerRow').forEach((el,i) => {
-										var prev_i = el.dataset.id;
-
-										if ( el == this.draggedEl && prev_i != i ) {
-											this.templates.splice(i, 0, this.templates.splice(prev_i, 1)[0]);
-										}
-										el.dataset.id = i;
-									});
-									this.store();
-								},
-								ondragover: (e) => {
-									e.preventDefault();
-									if ( e.currentTarget == this.draggedEl )
-										return;
-
-									let rect = e.currentTarget.getBoundingClientRect();
-									if (e.clientY > rect.top + rect.height / 2) {
-										e.currentTarget.parentNode.insertBefore(this.draggedEl, e.currentTarget.nextSibling);
-									} else {
-										e.currentTarget.parentNode.insertBefore(this.draggedEl, e.currentTarget);
-									}
-								}
-							},
-							[
-								$el(
-									"label",
-									{
-										textContent: "Name: ",
-										style: {
-											cursor: "grab",
-										},
-										onmousedown: (e) => {
-											// enable dragging only from the label
-											if (e.target.localName == 'label')
-												e.currentTarget.parentNode.draggable = 'true';
-										}
-									},
-									[
-										$el("input", {
-											value: t.name,
-											dataset: { name: t.name },
-											style: {
-												transitionProperty: 'background-color',
-												transitionDuration: '0s',
-											},
-											onchange: (e) => {
-												clearTimeout(this.saveVisualCue);
-												var el = e.target;
-												var row = el.parentNode.parentNode;
-												this.templates[row.dataset.id].name = el.value.trim() || 'untitled';
-												this.store();
-												el.style.backgroundColor = 'rgb(40, 95, 40)';
-												el.style.transitionDuration = '0s';
-												this.saveVisualCue = setTimeout(function () {
-													el.style.transitionDuration = '.7s';
-													el.style.backgroundColor = 'var(--comfy-input-bg)';
-												}, 15);
-											},
-											onkeypress: (e) => {
-												var el = e.target;
-												clearTimeout(this.saveVisualCue);
-												el.style.transitionDuration = '0s';
-												el.style.backgroundColor = 'var(--comfy-input-bg)';
-											},
-											$: (el) => (nameInput = el),
-										})
-									]
-								),
-								$el(
-									"div",
-									{},
-									[
-										$el("button", {
-											textContent: "Export",
-											style: {
-												fontSize: "12px",
-												fontWeight: "normal",
-											},
-											onclick: (e) => {
-												const json = JSON.stringify({templates: [t]}, null, 2); // convert the data to a JSON string
-												const blob = new Blob([json], {type: "application/json"});
-												const url = URL.createObjectURL(blob);
-												const a = $el("a", {
-													href: url,
-													download: (nameInput.value || t.name) + ".json",
-													style: {display: "none"},
-													parent: document.body,
-												});
-												a.click();
-												setTimeout(function () {
-													a.remove();
-													window.URL.revokeObjectURL(url);
-												}, 0);
-											},
-										}),
-										$el("button", {
-											textContent: "Delete",
-											style: {
-												fontSize: "12px",
-												color: "red",
-												fontWeight: "normal",
-											},
-											onclick: (e) => {
-												const item = e.target.parentNode.parentNode;
-												item.parentNode.removeChild(item);
-												this.templates.splice(item.dataset.id*1, 1);
-												this.store();
-												// update the rows index, setTimeout ensures that the list is updated
-												var that = this;
-												setTimeout(function (){
-													that.element.querySelectorAll('.tempateManagerRow').forEach((el,i) => {
-														el.dataset.id = i;
-													});
-												}, 0);
-											},
-										}),
-									]
-								),
-							]
-						)
-					];
-				})
-			)
-		);
-	}
-}
-
-app.registerExtension({
-	name: id,
-	setup() {
-		const manage = new ManageTemplates();
-
-		const clipboardAction = async (cb) => {
-			// We use the clipboard functions but dont want to overwrite the current user clipboard
-			// Restore it after we've run our callback
-			const old = localStorage.getItem("litegrapheditor_clipboard");
-			await cb();
-			localStorage.setItem("litegrapheditor_clipboard", old);
-		};
-
-		const orig = LGraphCanvas.prototype.getCanvasMenuOptions;
-		LGraphCanvas.prototype.getCanvasMenuOptions = function () {
-			const options = orig.apply(this, arguments);
-
-			options.push(null);
-			options.push({
-				content: `Save Selected as Template`,
-				disabled: !Object.keys(app.canvas.selected_nodes || {}).length,
-				callback: () => {
-					const name = prompt("Enter name");
-					if (!name?.trim()) return;
-
-					clipboardAction(() => {
-						app.canvas.copyToClipboard();
-						let data = localStorage.getItem("litegrapheditor_clipboard");
-						data = JSON.parse(data);
-						const nodeIds = Object.keys(app.canvas.selected_nodes);
-						for (let i = 0; i < nodeIds.length; i++) {
-							const node = app.graph.getNodeById(nodeIds[i]);
-							const nodeData = node?.constructor.nodeData;
-							
-							let groupData = GroupNodeHandler.getGroupData(node);
-							if (groupData) {
-								groupData = groupData.nodeData;
-								if (!data.groupNodes) {
-									data.groupNodes = {};
-								}
-								data.groupNodes[nodeData.name] = groupData;
-								data.nodes[i].type = nodeData.name;
-							}
-						}
-
-						manage.templates.push({
-							name,
-							data: JSON.stringify(data),
-						});
-						manage.store();
-					});
-				},
-			});
-
-			// Map each template to a menu item
-			const subItems = manage.templates.map((t) => {
-				return {
-					content: t.name,
-					callback: () => {
-						clipboardAction(async () => {
-							const data = JSON.parse(t.data);
-							await GroupNodeConfig.registerFromWorkflow(data.groupNodes, {});
-							localStorage.setItem("litegrapheditor_clipboard", t.data);
-							app.canvas.pasteFromClipboard();
-						});
-					},
-				};
-			});
-
-			subItems.push(null, {
-				content: "Manage",
-				callback: () => manage.show(),
-			});
-
-			options.push({
-				content: "Node Templates",
-				submenu: {
-					options: subItems,
-				},
-			});
-
-			return options;
-		};
-	},
-});

web/extensions/core/noteNode.js

@@ -1,41 +0,0 @@
-import {app} from "../../scripts/app.js";
-import {ComfyWidgets} from "../../scripts/widgets.js";
-// Node that add notes to your project
-
-app.registerExtension({
-    name: "Comfy.NoteNode",
-    registerCustomNodes() {
-        class NoteNode {
-            color=LGraphCanvas.node_colors.yellow.color;
-            bgcolor=LGraphCanvas.node_colors.yellow.bgcolor;
-            groupcolor = LGraphCanvas.node_colors.yellow.groupcolor;
-            constructor() {
-                if (!this.properties) {
-                    this.properties = {};
-                    this.properties.text="";
-                }
-
-                ComfyWidgets.STRING(this, "", ["", {default:this.properties.text, multiline: true}], app)
-
-                this.serialize_widgets = true;
-                this.isVirtualNode = true;
-
-            }
-
-
-        }
-
-        // Load default visibility
-
-        LiteGraph.registerNodeType(
-            "Note",
-            Object.assign(NoteNode, {
-                title_mode: LiteGraph.NORMAL_TITLE,
-                title: "Note",
-                collapsable: true,
-            })
-        );
-
-        NoteNode.category = "utils";
-    },
-});

web/extensions/core/rerouteNode.js

@@ -1,274 +0,0 @@
-import { app } from "../../scripts/app.js";
-import { mergeIfValid, getWidgetConfig, setWidgetConfig } from "./widgetInputs.js";
-
-// Node that allows you to redirect connections for cleaner graphs
-
-app.registerExtension({
-	name: "Comfy.RerouteNode",
-	registerCustomNodes(app) {
-		class RerouteNode {
-			constructor() {
-				if (!this.properties) {
-					this.properties = {};
-				}
-				this.properties.showOutputText = RerouteNode.defaultVisibility;
-				this.properties.horizontal = false;
-
-				this.addInput("", "*");
-				this.addOutput(this.properties.showOutputText ? "*" : "", "*");
-
-				this.onAfterGraphConfigured = function () {
-					requestAnimationFrame(() => {
-						this.onConnectionsChange(LiteGraph.INPUT, null, true, null);
-					});
-				};
-
-				this.onConnectionsChange = function (type, index, connected, link_info) {
-					this.applyOrientation();
-
-					// Prevent multiple connections to different types when we have no input
-					if (connected && type === LiteGraph.OUTPUT) {
-						// Ignore wildcard nodes as these will be updated to real types
-						const types = new Set(this.outputs[0].links.map((l) => app.graph.links[l].type).filter((t) => t !== "*"));
-						if (types.size > 1) {
-							const linksToDisconnect = [];
-							for (let i = 0; i < this.outputs[0].links.length - 1; i++) {
-								const linkId = this.outputs[0].links[i];
-								const link = app.graph.links[linkId];
-								linksToDisconnect.push(link);
-							}
-							for (const link of linksToDisconnect) {
-								const node = app.graph.getNodeById(link.target_id);
-								node.disconnectInput(link.target_slot);
-							}
-						}
-					}
-
-					// Find root input
-					let currentNode = this;
-					let updateNodes = [];
-					let inputType = null;
-					let inputNode = null;
-					while (currentNode) {
-						updateNodes.unshift(currentNode);
-						const linkId = currentNode.inputs[0].link;
-						if (linkId !== null) {
-							const link = app.graph.links[linkId];
-							if (!link) return;
-							const node = app.graph.getNodeById(link.origin_id);
-							const type = node.constructor.type;
-							if (type === "Reroute") {
-								if (node === this) {
-									// We've found a circle
-									currentNode.disconnectInput(link.target_slot);
-									currentNode = null;
-								} else {
-									// Move the previous node
-									currentNode = node;
-								}
-							} else {
-								// We've found the end
-								inputNode = currentNode;
-								inputType = node.outputs[link.origin_slot]?.type ?? null;
-								break;
-							}
-						} else {
-							// This path has no input node
-							currentNode = null;
-							break;
-						}
-					}
-
-					// Find all outputs
-					const nodes = [this];
-					let outputType = null;
-					while (nodes.length) {
-						currentNode = nodes.pop();
-						const outputs = (currentNode.outputs ? currentNode.outputs[0].links : []) || [];
-						if (outputs.length) {
-							for (const linkId of outputs) {
-								const link = app.graph.links[linkId];
-
-								// When disconnecting sometimes the link is still registered
-								if (!link) continue;
-
-								const node = app.graph.getNodeById(link.target_id);
-								const type = node.constructor.type;
-
-								if (type === "Reroute") {
-									// Follow reroute nodes
-									nodes.push(node);
-									updateNodes.push(node);
-								} else {
-									// We've found an output
-									const nodeOutType =
-										node.inputs && node.inputs[link?.target_slot] && node.inputs[link.target_slot].type
-											? node.inputs[link.target_slot].type
-											: null;
-									if (inputType && inputType !== "*" && nodeOutType !== inputType) {
-										// The output doesnt match our input so disconnect it
-										node.disconnectInput(link.target_slot);
-									} else {
-										outputType = nodeOutType;
-									}
-								}
-							}
-						} else {
-							// No more outputs for this path
-						}
-					}
-
-					const displayType = inputType || outputType || "*";
-					const color = LGraphCanvas.link_type_colors[displayType];
-
-					let widgetConfig;
-					let targetWidget;
-					let widgetType;
-					// Update the types of each node
-					for (const node of updateNodes) {
-						// If we dont have an input type we are always wildcard but we'll show the output type
-						// This lets you change the output link to a different type and all nodes will update
-						node.outputs[0].type = inputType || "*";
-						node.__outputType = displayType;
-						node.outputs[0].name = node.properties.showOutputText ? displayType : "";
-						node.size = node.computeSize();
-						node.applyOrientation();
-
-						for (const l of node.outputs[0].links || []) {
-							const link = app.graph.links[l];
-							if (link) {
-								link.color = color;
-
-								if (app.configuringGraph) continue;
-								const targetNode = app.graph.getNodeById(link.target_id);
-								const targetInput = targetNode.inputs?.[link.target_slot];
-								if (targetInput?.widget) {
-									const config = getWidgetConfig(targetInput);
-									if (!widgetConfig) {
-										widgetConfig = config[1] ?? {};
-										widgetType = config[0];
-									}
-									if (!targetWidget) {
-										targetWidget = targetNode.widgets?.find((w) => w.name === targetInput.widget.name);
-									}
-
-									const merged = mergeIfValid(targetInput, [config[0], widgetConfig]);
-									if (merged.customConfig) {
-										widgetConfig = merged.customConfig;
-									}
-								}
-							}
-						}
-					}
-
-					for (const node of updateNodes) {
-						if (widgetConfig && outputType) {
-							node.inputs[0].widget = { name: "value" };
-							setWidgetConfig(node.inputs[0], [widgetType ?? displayType, widgetConfig], targetWidget);
-						} else {
-							setWidgetConfig(node.inputs[0], null);
-						}
-					}
-
-					if (inputNode) {
-						const link = app.graph.links[inputNode.inputs[0].link];
-						if (link) {
-							link.color = color;
-						}
-					}
-				};
-
-				this.clone = function () {
-					const cloned = RerouteNode.prototype.clone.apply(this);
-					cloned.removeOutput(0);
-					cloned.addOutput(this.properties.showOutputText ? "*" : "", "*");
-					cloned.size = cloned.computeSize();
-					return cloned;
-				};
-
-				// This node is purely frontend and does not impact the resulting prompt so should not be serialized
-				this.isVirtualNode = true;
-			}
-
-			getExtraMenuOptions(_, options) {
-				options.unshift(
-					{
-						content: (this.properties.showOutputText ? "Hide" : "Show") + " Type",
-						callback: () => {
-							this.properties.showOutputText = !this.properties.showOutputText;
-							if (this.properties.showOutputText) {
-								this.outputs[0].name = this.__outputType || this.outputs[0].type;
-							} else {
-								this.outputs[0].name = "";
-							}
-							this.size = this.computeSize();
-							this.applyOrientation();
-							app.graph.setDirtyCanvas(true, true);
-						},
-					},
-					{
-						content: (RerouteNode.defaultVisibility ? "Hide" : "Show") + " Type By Default",
-						callback: () => {
-							RerouteNode.setDefaultTextVisibility(!RerouteNode.defaultVisibility);
-						},
-					},
-					{
-						// naming is inverted with respect to LiteGraphNode.horizontal
-						// LiteGraphNode.horizontal == true means that
-						// each slot in the inputs and outputs are layed out horizontally,
-						// which is the opposite of the visual orientation of the inputs and outputs as a node
-						content: "Set " + (this.properties.horizontal ? "Horizontal" : "Vertical"),
-						callback: () => {
-							this.properties.horizontal = !this.properties.horizontal;
-							this.applyOrientation();
-						},
-					}
-				);
-			}
-			applyOrientation() {
-				this.horizontal = this.properties.horizontal;
-				if (this.horizontal) {
-					// we correct the input position, because LiteGraphNode.horizontal
-					// doesn't account for title presence
-					// which reroute nodes don't have
-					this.inputs[0].pos = [this.size[0] / 2, 0];
-				} else {
-					delete this.inputs[0].pos;
-				}
-				app.graph.setDirtyCanvas(true, true);
-			}
-
-			computeSize() {
-				return [
-					this.properties.showOutputText && this.outputs && this.outputs.length
-						? Math.max(75, LiteGraph.NODE_TEXT_SIZE * this.outputs[0].name.length * 0.6 + 40)
-						: 75,
-					26,
-				];
-			}
-
-			static setDefaultTextVisibility(visible) {
-				RerouteNode.defaultVisibility = visible;
-				if (visible) {
-					localStorage["Comfy.RerouteNode.DefaultVisibility"] = "true";
-				} else {
-					delete localStorage["Comfy.RerouteNode.DefaultVisibility"];
-				}
-			}
-		}
-
-		// Load default visibility
-		RerouteNode.setDefaultTextVisibility(!!localStorage["Comfy.RerouteNode.DefaultVisibility"]);
-
-		LiteGraph.registerNodeType(
-			"Reroute",
-			Object.assign(RerouteNode, {
-				title_mode: LiteGraph.NO_TITLE,
-				title: "Reroute",
-				collapsable: false,
-			})
-		);
-
-		RerouteNode.category = "utils";
-	},
-});

web/extensions/core/saveImageExtraOutput.js

@@ -1,35 +0,0 @@
-import { app } from "../../scripts/app.js";
-import { applyTextReplacements } from "../../scripts/utils.js";
-// Use widget values and dates in output filenames
-
-app.registerExtension({
-	name: "Comfy.SaveImageExtraOutput",
-	async beforeRegisterNodeDef(nodeType, nodeData, app) {
-		if (nodeData.name === "SaveImage") {
-			const onNodeCreated = nodeType.prototype.onNodeCreated;
-			// When the SaveImage node is created we want to override the serialization of the output name widget to run our S&R
-			nodeType.prototype.onNodeCreated = function () {
-				const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
-
-				const widget = this.widgets.find((w) => w.name === "filename_prefix");
-				widget.serializeValue = () => {
-					return applyTextReplacements(app, widget.value);
-				};
-
-				return r;
-			};
-		} else {
-			// When any other node is created add a property to alias the node
-			const onNodeCreated = nodeType.prototype.onNodeCreated;
-			nodeType.prototype.onNodeCreated = function () {
-				const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
-
-				if (!this.properties || !("Node name for S&R" in this.properties)) {
-					this.addProperty("Node name for S&R", this.constructor.type, "string");
-				}
-
-				return r;
-			};
-		}
-	},
-});