mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-13 01:17:00 +08:00
Merge branch 'master' into execution_model_inversion
This commit is contained in:
Jacob Segal
commit
519d08e81c
@ -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
|
||||
|
||||
8
.ci/update_windows/update_comfyui_stable.bat
Executable file
8
.ci/update_windows/update_comfyui_stable.bat
Executable file
@ -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
|
||||
16
.github/workflows/pullrequest-ci-run.yml
vendored
16
.github/workflows/pullrequest-ci-run.yml
vendored
@ -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'
|
||||
})
|
||||
|
||||
@ -1,10 +1,4 @@
|
||||
# This is a temporary action during frontend TS migration.
|
||||
# This file should be removed after TS migration is completed.
|
||||
# The browser test is here to ensure TS repo is working the same way as the
|
||||
# current JS code.
|
||||
# If you are adding UI feature, please sync your changes to the TS repo:
|
||||
# huchenlei/ComfyUI_frontend and update test expectation files accordingly.
|
||||
name: Playwright Browser Tests CI
|
||||
name: Test server launches without errors
|
||||
|
||||
on:
|
||||
push:
|
||||
@ -21,15 +15,6 @@ jobs:
|
||||
with:
|
||||
repository: "comfyanonymous/ComfyUI"
|
||||
path: "ComfyUI"
|
||||
- name: Checkout ComfyUI_frontend
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
repository: "huchenlei/ComfyUI_frontend"
|
||||
path: "ComfyUI_frontend"
|
||||
ref: "fcc54d803e5b6a9b08a462a1d94899318c96dcbb"
|
||||
- uses: actions/setup-node@v3
|
||||
with:
|
||||
node-version: lts/*
|
||||
- uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: '3.8'
|
||||
@ -45,16 +30,6 @@ jobs:
|
||||
python main.py --cpu 2>&1 | tee console_output.log &
|
||||
wait-for-it --service 127.0.0.1:8188 -t 600
|
||||
working-directory: ComfyUI
|
||||
- name: Install ComfyUI_frontend dependencies
|
||||
run: |
|
||||
npm ci
|
||||
working-directory: ComfyUI_frontend
|
||||
- name: Install Playwright Browsers
|
||||
run: npx playwright install --with-deps
|
||||
working-directory: ComfyUI_frontend
|
||||
- name: Run Playwright tests
|
||||
run: npx playwright test
|
||||
working-directory: ComfyUI_frontend
|
||||
- name: Check for unhandled exceptions in server log
|
||||
run: |
|
||||
if grep -qE "Exception|Error" console_output.log; then
|
||||
@ -62,12 +37,6 @@ jobs:
|
||||
exit 1
|
||||
fi
|
||||
working-directory: ComfyUI
|
||||
- uses: actions/upload-artifact@v4
|
||||
if: always()
|
||||
with:
|
||||
name: playwright-report
|
||||
path: ComfyUI_frontend/playwright-report/
|
||||
retention-days: 30
|
||||
- uses: actions/upload-artifact@v4
|
||||
if: always()
|
||||
with:
|
||||
30
.github/workflows/test-ui.yaml
vendored
30
.github/workflows/test-ui.yaml
vendored
@ -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
|
||||
1
.gitignore
vendored
1
.gitignore
vendored
@ -19,3 +19,4 @@ venv/
|
||||
/user/
|
||||
*.log
|
||||
web_custom_versions/
|
||||
.DS_Store
|
||||
|
||||
@ -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,7 +501,10 @@ def load_controlnet(ckpt_path, model=None):
|
||||
logging.warning("leftover keys: {}".format(leftover_keys))
|
||||
controlnet_data = new_sd
|
||||
elif "controlnet_blocks.0.weight" in controlnet_data: #SD3 diffusers format
|
||||
return load_controlnet_mmdit(controlnet_data)
|
||||
if "double_blocks.0.img_attn.norm.key_norm.scale" in controlnet_data:
|
||||
return load_controlnet_flux_xlabs(controlnet_data)
|
||||
else:
|
||||
return load_controlnet_mmdit(controlnet_data)
|
||||
|
||||
pth_key = 'control_model.zero_convs.0.0.weight'
|
||||
pth = False
|
||||
|
||||
@ -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:
|
||||
|
||||
104
comfy/ldm/flux/controlnet_xlabs.py
Normal file
104
comfy/ldm/flux/controlnet_xlabs.py
Normal file
@ -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)
|
||||
@ -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_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
|
||||
img = img + img_mod1.gate * self.img_attn.proj(img_attn)
|
||||
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)
|
||||
|
||||
@ -38,7 +38,7 @@ class Flux(nn.Module):
|
||||
Transformer model for flow matching on sequences.
|
||||
"""
|
||||
|
||||
def __init__(self, image_model=None, dtype=None, device=None, operations=None, **kwargs):
|
||||
def __init__(self, image_model=None, final_layer=True, dtype=None, device=None, operations=None, **kwargs):
|
||||
super().__init__()
|
||||
self.dtype = dtype
|
||||
params = FluxParams(**kwargs)
|
||||
@ -83,7 +83,8 @@ class Flux(nn.Module):
|
||||
]
|
||||
)
|
||||
|
||||
self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations)
|
||||
if final_layer:
|
||||
self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device, operations=operations)
|
||||
|
||||
def forward_orig(
|
||||
self,
|
||||
@ -94,6 +95,7 @@ class Flux(nn.Module):
|
||||
timesteps: Tensor,
|
||||
y: Tensor,
|
||||
guidance: Tensor = None,
|
||||
control=None,
|
||||
) -> Tensor:
|
||||
if img.ndim != 3 or txt.ndim != 3:
|
||||
raise ValueError("Input img and txt tensors must have 3 dimensions.")
|
||||
@ -112,8 +114,15 @@ class Flux(nn.Module):
|
||||
ids = torch.cat((txt_ids, img_ids), dim=1)
|
||||
pe = self.pe_embedder(ids)
|
||||
|
||||
for block in self.double_blocks:
|
||||
img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
|
||||
for i in range(len(self.double_blocks)):
|
||||
img, txt = self.double_blocks[i](img=img, txt=txt, vec=vec, pe=pe)
|
||||
|
||||
if control is not None: #Controlnet
|
||||
control_o = control.get("output")
|
||||
if i < len(control_o):
|
||||
add = control_o[i]
|
||||
if add is not None:
|
||||
img += add
|
||||
|
||||
img = torch.cat((txt, img), 1)
|
||||
for block in self.single_blocks:
|
||||
@ -123,7 +132,7 @@ class Flux(nn.Module):
|
||||
img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
|
||||
return img
|
||||
|
||||
def forward(self, x, timestep, context, y, guidance, **kwargs):
|
||||
def forward(self, x, timestep, context, y, guidance, control=None, **kwargs):
|
||||
bs, c, h, w = x.shape
|
||||
patch_size = 2
|
||||
x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
|
||||
@ -138,5 +147,5 @@ class Flux(nn.Module):
|
||||
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
|
||||
|
||||
txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
|
||||
out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance)
|
||||
out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control)
|
||||
return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w]
|
||||
|
||||
@ -47,7 +47,7 @@ def reshape_for_broadcast(freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]], x
|
||||
|
||||
|
||||
def rotate_half(x):
|
||||
x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
|
||||
x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, S, H, D//2]
|
||||
return torch.stack([-x_imag, x_real], dim=-1).flatten(3)
|
||||
|
||||
|
||||
@ -78,10 +78,9 @@ def apply_rotary_emb(
|
||||
xk_out = None
|
||||
if isinstance(freqs_cis, tuple):
|
||||
cos, sin = reshape_for_broadcast(freqs_cis, xq, head_first) # [S, D]
|
||||
cos, sin = cos.to(xq.device), sin.to(xq.device)
|
||||
xq_out = (xq.float() * cos + rotate_half(xq.float()) * sin).type_as(xq)
|
||||
xq_out = (xq * cos + rotate_half(xq) * sin)
|
||||
if xk is not None:
|
||||
xk_out = (xk.float() * cos + rotate_half(xk.float()) * sin).type_as(xk)
|
||||
xk_out = (xk * cos + rotate_half(xk) * sin)
|
||||
else:
|
||||
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # [B, S, H, D//2]
|
||||
freqs_cis = reshape_for_broadcast(freqs_cis, xq_, head_first).to(xq.device) # [S, D//2] --> [1, S, 1, D//2]
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -1,3 +1,21 @@
|
||||
"""
|
||||
This file is part of ComfyUI.
|
||||
Copyright (C) 2024 Comfy
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import logging
|
||||
from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
|
||||
@ -77,10 +95,13 @@ class BaseModel(torch.nn.Module):
|
||||
self.device = device
|
||||
|
||||
if not unet_config.get("disable_unet_model_creation", False):
|
||||
if self.manual_cast_dtype is not None:
|
||||
operations = comfy.ops.manual_cast
|
||||
if model_config.custom_operations is None:
|
||||
if self.manual_cast_dtype is not None:
|
||||
operations = comfy.ops.manual_cast
|
||||
else:
|
||||
operations = comfy.ops.disable_weight_init
|
||||
else:
|
||||
operations = comfy.ops.disable_weight_init
|
||||
operations = model_config.custom_operations
|
||||
self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
|
||||
if comfy.model_management.force_channels_last():
|
||||
self.diffusion_model.to(memory_format=torch.channels_last)
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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,15 +308,21 @@ class LoadedModel:
|
||||
|
||||
load_weights = not self.weights_loaded
|
||||
|
||||
try:
|
||||
if lowvram_model_memory > 0 and load_weights:
|
||||
self.real_model = self.model.patch_model_lowvram(device_to=patch_model_to, lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights)
|
||||
else:
|
||||
self.real_model = self.model.patch_model(device_to=patch_model_to, patch_weights=load_weights)
|
||||
except Exception as e:
|
||||
self.model.unpatch_model(self.model.offload_device)
|
||||
self.model_unload()
|
||||
raise e
|
||||
if self.model.loaded_size() > 0:
|
||||
use_more_vram = lowvram_model_memory
|
||||
if use_more_vram == 0:
|
||||
use_more_vram = 1e32
|
||||
self.model_use_more_vram(use_more_vram)
|
||||
else:
|
||||
try:
|
||||
if lowvram_model_memory > 0 and load_weights:
|
||||
self.real_model = self.model.patch_model_lowvram(device_to=patch_model_to, lowvram_model_memory=lowvram_model_memory, force_patch_weights=force_patch_weights)
|
||||
else:
|
||||
self.real_model = self.model.patch_model(device_to=patch_model_to, patch_weights=load_weights)
|
||||
except Exception as e:
|
||||
self.model.unpatch_model(self.model.offload_device)
|
||||
self.model_unload()
|
||||
raise e
|
||||
|
||||
if is_intel_xpu() and not args.disable_ipex_optimize:
|
||||
self.real_model = ipex.optimize(self.real_model.eval(), graph_mode=True, concat_linear=True)
|
||||
@ -332,8 +338,9 @@ 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)
|
||||
return False
|
||||
freed = self.model.partially_unload(self.model.offload_device, memory_to_free)
|
||||
if freed >= memory_to_free:
|
||||
return False
|
||||
self.model.unpatch_model(self.model.offload_device, unpatch_weights=unpatch_weights)
|
||||
self.model.model_patches_to(self.model.offload_device)
|
||||
self.weights_loaded = self.weights_loaded and not unpatch_weights
|
||||
@ -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
|
||||
total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
|
||||
unload_model_clones(loaded_model.model, unload_weights_only=True, force_unload=False) #unload clones where the weights are different
|
||||
total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
|
||||
|
||||
for device in total_memory_required:
|
||||
if device != torch.device("cpu"):
|
||||
free_memory(total_memory_required[device] * 1.3 + extra_mem, device, models_already_loaded)
|
||||
for loaded_model in models_already_loaded:
|
||||
total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
|
||||
|
||||
for loaded_model in models_to_load:
|
||||
weights_unloaded = unload_model_clones(loaded_model.model, unload_weights_only=False, force_unload=False) #unload the rest of the clones where the weights can stay loaded
|
||||
if weights_unloaded is not None:
|
||||
loaded_model.weights_loaded = not weights_unloaded
|
||||
|
||||
for device in total_memory_required:
|
||||
if device != torch.device("cpu"):
|
||||
free_memory(total_memory_required[device] * 1.1 + extra_mem, device, models_already_loaded)
|
||||
|
||||
for loaded_model in models_to_load:
|
||||
model = loaded_model.model
|
||||
torch_dev = model.load_device
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
self.patch_weight_to_device(bias_key)
|
||||
weight_to = None
|
||||
if full_load:#TODO
|
||||
weight_to = device_to
|
||||
self.patch_weight_to_device(weight_key, device_to=weight_to) #TODO: speed this up without OOM
|
||||
self.patch_weight_to_device(bias_key, device_to=weight_to)
|
||||
m.to(device_to)
|
||||
logging.debug("lowvram: loaded module regularly {} {}".format(n, m))
|
||||
|
||||
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):
|
||||
|
||||
59
comfy/sd.py
59
comfy/sd.py
@ -62,7 +62,7 @@ def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
|
||||
|
||||
|
||||
class CLIP:
|
||||
def __init__(self, target=None, embedding_directory=None, no_init=False, tokenizer_data={}):
|
||||
def __init__(self, target=None, embedding_directory=None, no_init=False, tokenizer_data={}, parameters=0):
|
||||
if no_init:
|
||||
return
|
||||
params = target.params.copy()
|
||||
@ -71,20 +71,24 @@ class CLIP:
|
||||
|
||||
load_device = model_management.text_encoder_device()
|
||||
offload_device = model_management.text_encoder_offload_device()
|
||||
params['device'] = offload_device
|
||||
dtype = model_management.text_encoder_dtype(load_device)
|
||||
params['dtype'] = dtype
|
||||
|
||||
params['device'] = model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype))
|
||||
self.cond_stage_model = clip(**(params))
|
||||
|
||||
for dt in self.cond_stage_model.dtypes:
|
||||
if not model_management.supports_cast(load_device, dt):
|
||||
load_device = offload_device
|
||||
if params['device'] != offload_device:
|
||||
self.cond_stage_model.to(offload_device)
|
||||
logging.warning("Had to shift TE back.")
|
||||
|
||||
self.tokenizer = tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
|
||||
self.patcher = comfy.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
|
||||
if params['device'] == load_device:
|
||||
model_management.load_models_gpu([self.patcher], force_full_load=True)
|
||||
self.layer_idx = None
|
||||
logging.debug("CLIP model load device: {}, offload device: {}".format(load_device, offload_device))
|
||||
logging.debug("CLIP model load device: {}, offload device: {}, current: {}".format(load_device, offload_device, params['device']))
|
||||
|
||||
def clone(self):
|
||||
n = CLIP(no_init=True)
|
||||
@ -456,7 +460,11 @@ def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DI
|
||||
clip_target.clip = comfy.text_encoders.sd3_clip.SD3ClipModel
|
||||
clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
|
||||
|
||||
clip = CLIP(clip_target, embedding_directory=embedding_directory)
|
||||
parameters = 0
|
||||
for c in clip_data:
|
||||
parameters += comfy.utils.calculate_parameters(c)
|
||||
|
||||
clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters)
|
||||
for c in clip_data:
|
||||
m, u = clip.load_sd(c)
|
||||
if len(m) > 0:
|
||||
@ -498,15 +506,19 @@ def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_cl
|
||||
|
||||
return (model, clip, vae)
|
||||
|
||||
def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True):
|
||||
def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True, model_options={}):
|
||||
sd = comfy.utils.load_torch_file(ckpt_path)
|
||||
sd_keys = sd.keys()
|
||||
out = load_state_dict_guess_config(sd, output_vae, output_clip, output_clipvision, embedding_directory, output_model, model_options)
|
||||
if out is None:
|
||||
raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
|
||||
return out
|
||||
|
||||
def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True, model_options={}):
|
||||
clip = None
|
||||
clipvision = None
|
||||
vae = None
|
||||
model = None
|
||||
model_patcher = None
|
||||
clip_target = None
|
||||
|
||||
diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
|
||||
parameters = comfy.utils.calculate_parameters(sd, diffusion_model_prefix)
|
||||
@ -515,13 +527,18 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
|
||||
|
||||
model_config = model_detection.model_config_from_unet(sd, diffusion_model_prefix)
|
||||
if model_config is None:
|
||||
raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
|
||||
return None
|
||||
|
||||
unet_weight_dtype = list(model_config.supported_inference_dtypes)
|
||||
if weight_dtype is not None:
|
||||
unet_weight_dtype.append(weight_dtype)
|
||||
|
||||
unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype)
|
||||
model_config.custom_operations = model_options.get("custom_operations", None)
|
||||
unet_dtype = model_options.get("weight_dtype", None)
|
||||
|
||||
if unet_dtype is None:
|
||||
unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype)
|
||||
|
||||
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
|
||||
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
|
||||
|
||||
@ -545,7 +562,8 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
|
||||
if clip_target is not None:
|
||||
clip_sd = model_config.process_clip_state_dict(sd)
|
||||
if len(clip_sd) > 0:
|
||||
clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd)
|
||||
parameters = comfy.utils.calculate_parameters(clip_sd)
|
||||
clip = CLIP(clip_target, embedding_directory=embedding_directory, tokenizer_data=clip_sd, parameters=parameters)
|
||||
m, u = clip.load_sd(clip_sd, full_model=True)
|
||||
if len(m) > 0:
|
||||
m_filter = list(filter(lambda a: ".logit_scale" not in a and ".transformer.text_projection.weight" not in a, m))
|
||||
@ -567,12 +585,13 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
|
||||
model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device())
|
||||
if inital_load_device != torch.device("cpu"):
|
||||
logging.info("loaded straight to GPU")
|
||||
model_management.load_model_gpu(model_patcher)
|
||||
model_management.load_models_gpu([model_patcher], force_full_load=True)
|
||||
|
||||
return (model_patcher, clip, vae, clipvision)
|
||||
|
||||
|
||||
def load_unet_state_dict(sd, dtype=None): #load unet in diffusers or regular format
|
||||
def load_diffusion_model_state_dict(sd, model_options={}): #load unet in diffusers or regular format
|
||||
dtype = model_options.get("dtype", None)
|
||||
|
||||
#Allow loading unets from checkpoint files
|
||||
diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
|
||||
@ -614,6 +633,7 @@ def load_unet_state_dict(sd, dtype=None): #load unet in diffusers or regular for
|
||||
|
||||
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
|
||||
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
|
||||
model_config.custom_operations = model_options.get("custom_operations", None)
|
||||
model = model_config.get_model(new_sd, "")
|
||||
model = model.to(offload_device)
|
||||
model.load_model_weights(new_sd, "")
|
||||
@ -622,14 +642,23 @@ def load_unet_state_dict(sd, dtype=None): #load unet in diffusers or regular for
|
||||
logging.info("left over keys in unet: {}".format(left_over))
|
||||
return comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device)
|
||||
|
||||
def load_unet(unet_path, dtype=None):
|
||||
|
||||
def load_diffusion_model(unet_path, model_options={}):
|
||||
sd = comfy.utils.load_torch_file(unet_path)
|
||||
model = load_unet_state_dict(sd, dtype=dtype)
|
||||
model = load_diffusion_model_state_dict(sd, model_options=model_options)
|
||||
if model is None:
|
||||
logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path))
|
||||
raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
|
||||
return model
|
||||
|
||||
def load_unet(unet_path, dtype=None):
|
||||
print("WARNING: the load_unet function has been deprecated and will be removed please switch to: load_diffusion_model")
|
||||
return load_diffusion_model(unet_path, model_options={"dtype": dtype})
|
||||
|
||||
def load_unet_state_dict(sd, dtype=None):
|
||||
print("WARNING: the load_unet_state_dict function has been deprecated and will be removed please switch to: load_diffusion_model_state_dict")
|
||||
return load_diffusion_model_state_dict(sd, model_options={"dtype": dtype})
|
||||
|
||||
def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, metadata=None, extra_keys={}):
|
||||
clip_sd = None
|
||||
load_models = [model]
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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",
|
||||
"attn.to_out.0.bias": "img_attn.proj.bias",
|
||||
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:
|
||||
|
||||
1
main.py
1
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")
|
||||
|
||||
2
model_filemanager/__init__.py
Normal file
2
model_filemanager/__init__.py
Normal file
@ -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
|
||||
240
model_filemanager/download_models.py
Normal file
240
model_filemanager/download_models.py
Normal file
@ -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
|
||||
111
nodes.py
111
nodes.py
@ -47,11 +47,18 @@ MAX_RESOLUTION=16384
|
||||
class CLIPTextEncode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"text": ("STRING", {"multiline": True, "dynamicPrompts": True}), "clip": ("CLIP", )}}
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"multiline": True, "dynamicPrompts": True, "tooltip": "The text to be encoded."}),
|
||||
"clip": ("CLIP", {"tooltip": "The CLIP model used for encoding the text."})
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("CONDITIONING",)
|
||||
OUTPUT_TOOLTIPS = ("A conditioning containing the embedded text used to guide the diffusion model.",)
|
||||
FUNCTION = "encode"
|
||||
|
||||
CATEGORY = "conditioning"
|
||||
DESCRIPTION = "Encodes a text prompt using a CLIP model into an embedding that can be used to guide the diffusion model towards generating specific images."
|
||||
|
||||
def encode(self, clip, text):
|
||||
tokens = clip.tokenize(text)
|
||||
@ -260,11 +267,18 @@ class ConditioningSetTimestepRange:
|
||||
class VAEDecode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "samples": ("LATENT", ), "vae": ("VAE", )}}
|
||||
return {
|
||||
"required": {
|
||||
"samples": ("LATENT", {"tooltip": "The latent to be decoded."}),
|
||||
"vae": ("VAE", {"tooltip": "The VAE model used for decoding the latent."})
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
OUTPUT_TOOLTIPS = ("The decoded image.",)
|
||||
FUNCTION = "decode"
|
||||
|
||||
CATEGORY = "latent"
|
||||
DESCRIPTION = "Decodes latent images back into pixel space images."
|
||||
|
||||
def decode(self, vae, samples):
|
||||
return (vae.decode(samples["samples"]), )
|
||||
@ -506,12 +520,19 @@ class CheckpointLoader:
|
||||
class CheckpointLoaderSimple:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ),
|
||||
}}
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), {"tooltip": "The name of the checkpoint (model) to load."}),
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
OUTPUT_TOOLTIPS = ("The model used for denoising latents.",
|
||||
"The CLIP model used for encoding text prompts.",
|
||||
"The VAE model used for encoding and decoding images to and from latent space.")
|
||||
FUNCTION = "load_checkpoint"
|
||||
|
||||
CATEGORY = "loaders"
|
||||
DESCRIPTION = "Loads a diffusion model checkpoint, diffusion models are used to denoise latents."
|
||||
|
||||
def load_checkpoint(self, ckpt_name):
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
@ -582,16 +603,22 @@ class LoraLoader:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "model": ("MODEL",),
|
||||
"clip": ("CLIP", ),
|
||||
"lora_name": (folder_paths.get_filename_list("loras"), ),
|
||||
"strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
|
||||
"strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
|
||||
}}
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}),
|
||||
"clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}),
|
||||
"lora_name": (folder_paths.get_filename_list("loras"), {"tooltip": "The name of the LoRA."}),
|
||||
"strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the diffusion model. This value can be negative."}),
|
||||
"strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP")
|
||||
OUTPUT_TOOLTIPS = ("The modified diffusion model.", "The modified CLIP model.")
|
||||
FUNCTION = "load_lora"
|
||||
|
||||
CATEGORY = "loaders"
|
||||
DESCRIPTION = "LoRAs are used to modify diffusion and CLIP models, altering the way in which latents are denoised such as applying styles. Multiple LoRA nodes can be linked together."
|
||||
|
||||
def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
|
||||
if strength_model == 0 and strength_clip == 0:
|
||||
@ -826,14 +853,14 @@ class UNETLoader:
|
||||
CATEGORY = "advanced/loaders"
|
||||
|
||||
def load_unet(self, unet_name, weight_dtype):
|
||||
dtype = None
|
||||
model_options = {}
|
||||
if weight_dtype == "fp8_e4m3fn":
|
||||
dtype = torch.float8_e4m3fn
|
||||
model_options["dtype"] = torch.float8_e4m3fn
|
||||
elif weight_dtype == "fp8_e5m2":
|
||||
dtype = torch.float8_e5m2
|
||||
model_options["dtype"] = torch.float8_e5m2
|
||||
|
||||
unet_path = folder_paths.get_full_path("unet", unet_name)
|
||||
model = comfy.sd.load_unet(unet_path, dtype=dtype)
|
||||
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
|
||||
return (model,)
|
||||
|
||||
class CLIPLoader:
|
||||
@ -1033,13 +1060,19 @@ class EmptyLatentImage:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
|
||||
"height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}}
|
||||
return {
|
||||
"required": {
|
||||
"width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The width of the latent images in pixels."}),
|
||||
"height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8, "tooltip": "The height of the latent images in pixels."}),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."})
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
OUTPUT_TOOLTIPS = ("The empty latent image batch.",)
|
||||
FUNCTION = "generate"
|
||||
|
||||
CATEGORY = "latent"
|
||||
DESCRIPTION = "Create a new batch of empty latent images to be denoised via sampling."
|
||||
|
||||
def generate(self, width, height, batch_size=1):
|
||||
latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
|
||||
@ -1359,24 +1392,27 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
|
||||
class KSampler:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"model": ("MODEL",),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
||||
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
|
||||
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
|
||||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
|
||||
"positive": ("CONDITIONING", ),
|
||||
"negative": ("CONDITIONING", ),
|
||||
"latent_image": ("LATENT", ),
|
||||
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
}
|
||||
}
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}),
|
||||
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}),
|
||||
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."}),
|
||||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
|
||||
"positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}),
|
||||
"negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
|
||||
"latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}),
|
||||
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising applied, lower values will maintain the structure of the initial image allowing for image to image sampling."}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
OUTPUT_TOOLTIPS = ("The denoised latent.",)
|
||||
FUNCTION = "sample"
|
||||
|
||||
CATEGORY = "sampling"
|
||||
DESCRIPTION = "Uses the provided model, positive and negative conditioning to denoise the latent image."
|
||||
|
||||
def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
|
||||
return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
|
||||
@ -1424,11 +1460,15 @@ class SaveImage:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":
|
||||
{"images": ("IMAGE", ),
|
||||
"filename_prefix": ("STRING", {"default": "ComfyUI"})},
|
||||
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
||||
}
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE", {"tooltip": "The images to save."}),
|
||||
"filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "save_images"
|
||||
@ -1436,6 +1476,7 @@ class SaveImage:
|
||||
OUTPUT_NODE = True
|
||||
|
||||
CATEGORY = "image"
|
||||
DESCRIPTION = "Saves the input images to your ComfyUI output directory."
|
||||
|
||||
def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
|
||||
filename_prefix += self.prefix_append
|
||||
|
||||
38
server.py
38
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)
|
||||
|
||||
|
||||
1
tests-ui/.gitignore
vendored
1
tests-ui/.gitignore
vendored
@ -1 +0,0 @@
|
||||
node_modules
|
||||
@ -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);
|
||||
});
|
||||
@ -1,4 +0,0 @@
|
||||
{
|
||||
"presets": ["@babel/preset-env"],
|
||||
"plugins": ["babel-plugin-transform-import-meta"]
|
||||
}
|
||||
@ -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";
|
||||
};
|
||||
@ -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;
|
||||
5586
tests-ui/package-lock.json
generated
5586
tests-ui/package-lock.json
generated
File diff suppressed because it is too large
Load Diff
@ -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"
|
||||
}
|
||||
}
|
||||
@ -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();
|
||||
@ -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");
|
||||
});
|
||||
});
|
||||
File diff suppressed because it is too large
Load Diff
@ -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();
|
||||
});
|
||||
});
|
||||
});
|
||||
@ -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,
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
});
|
||||
@ -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 };
|
||||
},
|
||||
};
|
||||
@ -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];
|
||||
}
|
||||
@ -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 = "";
|
||||
}
|
||||
@ -1,6 +0,0 @@
|
||||
export const nop = new Proxy(function () {}, {
|
||||
get: () => nop,
|
||||
set: () => true,
|
||||
apply: () => nop,
|
||||
construct: () => nop,
|
||||
});
|
||||
@ -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;
|
||||
},
|
||||
}));
|
||||
}
|
||||
0
tests-unit/prompt_server_test/__init__.py
Normal file
0
tests-unit/prompt_server_test/__init__.py
Normal file
321
tests-unit/prompt_server_test/download_models_test.py
Normal file
321
tests-unit/prompt_server_test/download_models_test.py
Normal file
@ -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
|
||||
@ -1 +1,3 @@
|
||||
pytest>=7.8.0
|
||||
pytest-aiohttp
|
||||
pytest-asyncio
|
||||
|
||||
122
web/extensions/core/tooltips.js
Normal file
122
web/extensions/core/tooltips.js
Normal file
@ -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);
|
||||
}
|
||||
},
|
||||
});
|
||||
},
|
||||
});
|
||||
@ -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];
|
||||
|
||||
|
||||
@ -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();
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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);
|
||||
}
|
||||
|
||||
@ -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;
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user