Merge branch 'master' into execution_model_inversion

.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

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

@@ -35,3 +35,19 @@ jobs:
           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-launch.yml

@@ -1,10 +1,4 @@
-# This is a temporary action during frontend TS migration.
+name: Test server launches without errors
-# This file should be removed after TS migration is completed.
-# The browser test is here to ensure TS repo is working the same way as the
-# current JS code.
-# If you are adding UI feature, please sync your changes to the TS repo:
-# huchenlei/ComfyUI_frontend and update test expectation files accordingly.
-name: Playwright Browser Tests CI
 
 on:
   push:
@@ -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

@@ -19,3 +19,4 @@ venv/
 /user/
 *.log
 web_custom_versions/
+.DS_Store

comfy/controlnet.py

@@ -34,6 +34,8 @@ 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):
     current_batch_size = tensor.shape[0]
@@ -322,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)
@@ -415,6 +418,7 @@ def load_controlnet_mmdit(sd):
     control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
     return control
 
+
 def load_controlnet_hunyuandit(controlnet_data):
     model_config, operations, load_device, unet_dtype, manual_cast_dtype = controlnet_config(controlnet_data)
 
@@ -426,6 +430,15 @@ def load_controlnet_hunyuandit(controlnet_data):
     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
@@ -488,6 +501,9 @@ 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
+        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'

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

@@ -170,8 +170,8 @@ class DoubleStreamBlock(nn.Module):
         txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
 
         # calculate the img bloks
-        img += img_mod1.gate * self.img_attn.proj(img_attn)
+        img = img + img_mod1.gate * self.img_attn.proj(img_attn)
-        img += img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
+        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_mod1.gate * self.txt_attn.proj(txt_attn)

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,6 +83,7 @@ class Flux(nn.Module):
             ]
         )
 
+        if final_layer:
             self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations)
 
     def forward_orig(
@@ -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:
+        for i in range(len(self.double_blocks)):
-            img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+            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 * cos + rotate_half(xq) * sin)
-        xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq)
         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/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
 
 

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/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 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 = 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

@@ -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

@@ -296,7 +296,7 @@ class LoadedModel:
 
     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()
 
@@ -308,6 +308,12 @@ class LoadedModel:
 
         load_weights = not self.weights_loaded
 
+        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)
@@ -332,7 +338,8 @@ class LoadedModel:
     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():
-                self.model.partially_unload(self.model.offload_device, memory_to_free)
+                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)
@@ -428,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) + 100 * 1024 * 1024
+    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) + 100 * 1024 * 1024
+        minimum_memory_required = max(inference_memory, minimum_memory_required + 300 * 1024 * 1024)
 
     models = set(models)
 
@@ -484,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
+        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:
+    for loaded_model in models_already_loaded:
-        if device != torch.device("cpu"):
+        total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
-            free_memory(total_memory_required[device] * 1.3 + extra_mem, device, models_already_loaded)
 
     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
@@ -504,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()))
@@ -675,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

@@ -102,7 +102,7 @@ class ModelPatcher:
         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
@@ -355,13 +355,14 @@ class ModelPatcher:
 
         return self.model
 
-    def lowvram_load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False):
+    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
@@ -401,13 +402,16 @@ class ModelPatcher:
                         if weight.device == device_to:
                             continue
 
-                    self.patch_weight_to_device(weight_key) #TODO: speed this up without OOM
+                    weight_to = None
-                    self.patch_weight_to_device(bias_key)
+                    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))
 
         if lowvram_counter > 0:
-            logging.info("loaded in lowvram mode {}".format(lowvram_model_memory / (1024 * 1024)))
+            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)))
@@ -665,12 +669,15 @@ class ModelPatcher:
         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():
-            pass #TODO: Full load
+            full_load = True
         current_used = self.model.model_loaded_weight_memory
-        self.lowvram_load(device_to, lowvram_model_memory=current_used + extra_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):

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)
 
+    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/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

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

@@ -457,8 +457,27 @@ def flux_to_diffusers(mmdit_config, output_prefix=""):
             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",
+        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:
@@ -474,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:

main.py

@@ -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",),
+        return {
-                              "clip": ("CLIP", ),
+            "required": { 
-                              "lora_name": (folder_paths.get_filename_list("loras"), ),
+                "model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}),
-                              "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
+                "clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}),
-                              "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
+                "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}),
+        return {
-                              "height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
+            "required": { 
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}}
+                "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":
+        return {
-                    {"model": ("MODEL",),
+            "required": {
-                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}),
-                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}),
-                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}),
-                    "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
+                "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."}),
-                    "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
+                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."}),
-                    "positive": ("CONDITIONING", ),
+                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
-                    "negative": ("CONDITIONING", ),
+                "positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}),
-                    "latent_image": ("LATENT", ),
+                "negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
-                    "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "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,10 +1460,14 @@ class SaveImage:
 
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": 
+        return {
-                    {"images": ("IMAGE", ),
+            "required": {
-                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
+                "images": ("IMAGE", {"tooltip": "The images to save."}),
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+                "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 = ()
@@ -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):
         filename_prefix += self.prefix_append

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]
@@ -438,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")
@@ -561,6 +565,36 @@ class PromptServer():
 
             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)
 

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

web/extensions/core/tooltips.js

@@ -0,0 +1,122 @@
+import { app } from "../../scripts/app.js";
+import { $el } from "../../scripts/ui.js";
+
+// Adds support for tooltips
+
+function getHoveredWidget() {
+	if (!app) {
+		return;
+	}
+
+	const node = app.canvas.node_over;
+	if (!node.widgets) return;
+
+	const graphPos = app.canvas.graph_mouse;
+
+	const x = graphPos[0] - node.pos[0];
+	const y = graphPos[1] - node.pos[1];
+
+	for (const w of node.widgets) {
+		let widgetWidth, widgetHeight;
+		if (w.computeSize) {
+			const sz = w.computeSize();
+			widgetWidth = sz[0];
+			widgetHeight = sz[1];
+		} else {
+			widgetWidth = w.width || node.size[0];
+			widgetHeight = LiteGraph.NODE_WIDGET_HEIGHT;
+		}
+
+		if (w.last_y !== undefined && x >= 6 && x <= widgetWidth - 12 && y >= w.last_y && y <= w.last_y + widgetHeight) {
+			return w;
+		}
+	}
+}
+
+app.registerExtension({
+	name: "Comfy.Tooltips",
+	setup() {
+		const tooltipEl = $el("div.comfy-graph-tooltip", {
+			parent: document.body,
+		});
+		let idleTimeout;
+
+		const hideTooltip = () => {
+			tooltipEl.style.display = "none";
+		};
+		const showTooltip = (tooltip) => {
+			if (!tooltip) return;
+
+			tooltipEl.textContent = tooltip;
+			tooltipEl.style.display = "block";
+			tooltipEl.style.left = app.canvas.mouse[0] + "px";
+			tooltipEl.style.top = app.canvas.mouse[1] + "px";
+			const rect = tooltipEl.getBoundingClientRect();
+			if (rect.right > window.innerWidth) {
+				tooltipEl.style.left = app.canvas.mouse[0] - rect.width + "px";
+			}
+
+			if (rect.top < 0) {
+				tooltipEl.style.top = app.canvas.mouse[1] + rect.height + "px";
+			}
+		};
+		const getInputTooltip = (nodeData, name) => {
+			const inputDef = nodeData.input?.required?.[name] ?? nodeData.input?.optional?.[name];
+			return inputDef?.[1]?.tooltip;
+		};
+		const onIdle = () => {
+			const { canvas } = app;
+			const node = canvas.node_over;
+			if (!node) return;
+
+			const nodeData = node.constructor.nodeData ?? {};
+
+			if (node.constructor.title_mode !== LiteGraph.NO_TITLE && canvas.graph_mouse[1] < node.pos[1]) {
+				return showTooltip(nodeData.description);
+			}
+
+			if (node.flags?.collapsed) return;
+
+			const inputSlot = canvas.isOverNodeInput(node, canvas.graph_mouse[0], canvas.graph_mouse[1], [0, 0]);
+			if (inputSlot !== -1) {
+				const inputName = node.inputs[inputSlot].name;
+				return showTooltip(getInputTooltip(nodeData, inputName));
+			}
+
+			const outputSlot = canvas.isOverNodeOutput(node, canvas.graph_mouse[0], canvas.graph_mouse[1], [0, 0]);
+			if (outputSlot !== -1) {
+				return showTooltip(nodeData.output_tooltips?.[outputSlot]);
+			}
+
+			const widget = getHoveredWidget();
+			// Dont show for DOM widgets, these use native browser tooltips as we dont get proper mouse events on these
+			if (widget && !widget.element) {
+				return showTooltip(widget.tooltip ?? getInputTooltip(nodeData, widget.name));
+			}
+		};
+
+		const onMouseMove = (e) => {
+			hideTooltip();
+			clearTimeout(idleTimeout);
+			
+			if(e.target.nodeName !== "CANVAS") return
+			idleTimeout = setTimeout(onIdle, 500);
+		};
+
+		app.ui.settings.addSetting({
+			id: "Comfy.EnableTooltips",
+			name: "Enable Tooltips",
+			type: "boolean",
+			defaultValue: true,
+			onChange(value) {
+				if (value) {
+					window.addEventListener("mousemove", onMouseMove);
+					window.addEventListener("click", hideTooltip);
+				} else {
+					window.removeEventListener("mousemove", onMouseMove);
+					window.removeEventListener("click", hideTooltip);
+				}
+			},
+		});
+	},
+});

web/extensions/core/widgetInputs.js

@@ -181,7 +181,7 @@ export function mergeIfValid(output, config2, forceUpdate, recreateWidget, confi
 
 	const isNumber = config1[0] === "INT" || config1[0] === "FLOAT";
 	for (const k of keys.values()) {
-		if (k !== "default" && k !== "forceInput" && k !== "defaultInput" && k !== "control_after_generate" && k !== "multiline") {
+		if (k !== "default" && k !== "forceInput" && k !== "defaultInput" && k !== "control_after_generate" && k !== "multiline" && k !== "tooltip") {
 			let v1 = config1[1][k];
 			let v2 = config2[1]?.[k];
 

web/scripts/app.js

@@ -1713,9 +1713,10 @@ export class ComfyApp {
 				for (const o in nodeData["output"]) {
 					let output = nodeData["output"][o];
 					if(output instanceof Array) output = "COMBO";
+					const outputTooltip = nodeData["output_tooltips"]?.[o];
 					const outputName = nodeData["output_name"][o] || output;
 					const outputShape = nodeData["output_is_list"][o] ? LiteGraph.GRID_SHAPE : LiteGraph.CIRCLE_SHAPE ;
-					this.addOutput(outputName, output, { shape: outputShape });
+					this.addOutput(outputName, output, { shape: outputShape, tooltip: outputTooltip });
 				}
 
 				const s = this.computeSize();

web/scripts/domWidget.js

@@ -223,6 +223,12 @@ LGraphNode.prototype.addDOMWidget = function (name, type, element, options) {
 		document.addEventListener("mousedown", mouseDownHandler);
 	}
 
+	const { nodeData } = this.constructor;
+	const tooltip = (nodeData?.input.required?.[name] ?? nodeData?.input.optional?.[name])?.[1]?.tooltip;
+	if (tooltip && !element.title) {
+		element.title = tooltip;
+	}
+
 	const widget = {
 		type,
 		name,

web/scripts/widgets.js

@@ -75,6 +75,7 @@ export function addValueControlWidgets(node, targetWidget, defaultValue = "rando
 			serialize: false, // Don't include this in prompt.
 		}
 	);
+	valueControl.tooltip = "Allows the linked widget to be changed automatically, for example randomizing the noise seed.";
 	valueControl[IS_CONTROL_WIDGET] = true;
 	updateControlWidgetLabel(valueControl);
 	widgets.push(valueControl);
@@ -95,6 +96,7 @@ export function addValueControlWidgets(node, targetWidget, defaultValue = "rando
 			}
 		);
 		updateControlWidgetLabel(comboFilter);
+		comboFilter.tooltip = "Allows for filtering the list of values when changing the value via the control generate mode. Allows for RegEx matches in the format /abc/ to only filter to values containing 'abc'."
 
 		widgets.push(comboFilter);
 	}

web/style.css

@@ -645,3 +645,20 @@ dialog::backdrop {
 audio.comfy-audio.empty-audio-widget {
 	display: none;
 }
+
+.comfy-graph-tooltip {
+	background: var(--comfy-input-bg);
+	border-radius: 5px;
+	box-shadow: 0 0 5px rgba(0, 0, 0, 0.4);
+	color: var(--input-text);
+	display: none;
+	font-family: sans-serif;
+	left: 0;
+	max-width: 30vw;
+	padding: 4px 8px;
+	position: absolute;
+	top: 0;
+	transform: translate(5px, calc(-100% - 5px));
+	white-space: pre-wrap;
+	z-index: 99999;
+}