From 54fca4a218421b9938e0b09bdf9025d87b50bfae Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 9 Sep 2024 16:23:21 -0400 Subject: [PATCH 01/26] If host does not contain a port only compare the hostnames. --- server.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/server.py b/server.py index 93147f6c5..a1fca49cc 100644 --- a/server.py +++ b/server.py @@ -83,17 +83,22 @@ def create_cors_middleware(allowed_origin: str): def create_origin_only_middleware(): @web.middleware async def origin_only_middleware(request: web.Request, handler): + #this code is used to prevent the case where a random website can queue comfy workflows by making a POST to 127.0.0.1 which browsers don't prevent for some dumb reason. + #in that case the Host and Origin hostnames won't match + #I know the proper fix would be to add a cookie but this should take care of the problem in the meantime if 'Host' in request.headers and 'Origin' in request.headers: host = request.headers['Host'] origin = request.headers['Origin'] host_domain = host.lower() parsed = urllib.parse.urlparse(origin) origin_domain = parsed.netloc.lower() - if parsed.port is None: #if origin doesn't have a port strip it from the host to handle weird browsers - result = urllib.parse.urlsplit('//' + host_domain) - host_domain = result.hostname + host_domain_parsed = urllib.parse.urlsplit('//' + host_domain) + if parsed.port is None: #if origin doesn't have a port strip it from the host to handle weird browsers, same for host + host_domain = host_domain_parsed.hostname + if host_domain_parsed.port is None: + origin_domain = parsed.hostname - if len(host_domain) > 0 and len(origin_domain) > 0: + if host_domain is not None and origin_domain is not None and len(host_domain) > 0 and len(origin_domain) > 0: if host_domain != origin_domain: logging.warning("WARNING: request with non matching host and origin {} != {}, returning 403".format(host_domain, origin_domain)) return web.Response(status=403) From 9a7444e39fc62cdf694000f81f8f493815d73322 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 10 Sep 2024 00:21:33 -0400 Subject: [PATCH 02/26] Add diffusion_models to the extra_model_paths.yaml.example --- extra_model_paths.yaml.example | 3 +++ 1 file changed, 3 insertions(+) diff --git a/extra_model_paths.yaml.example b/extra_model_paths.yaml.example index 846d04dbe..d4fd55c87 100644 --- a/extra_model_paths.yaml.example +++ b/extra_model_paths.yaml.example @@ -30,6 +30,9 @@ a111: # clip_vision: models/clip_vision/ # configs: models/configs/ # controlnet: models/controlnet/ +# diffusion_models: | +# models/diffusion_models +# models/unet # embeddings: models/embeddings/ # loras: models/loras/ # upscale_models: models/upscale_models/ From 9fa8faa44a0226eb9a0980006cf62ca2c1c4973a Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Mon, 9 Sep 2024 21:33:44 -0700 Subject: [PATCH 03/26] Expand user directory for basepath in extra_models_paths.yaml (#4857) * Expand user path. * Add test. * Add unit test for expanding base path. * Simplify unit test. * Remove comment. * Remove comment. * Checkpoints. * Refactor. --- main.py | 27 ++--------- tests-unit/utils/extra_config_test.py | 69 +++++++++++++++++++++++++++ utils/__init__.py | 0 utils/extra_config.py | 25 ++++++++++ 4 files changed, 97 insertions(+), 24 deletions(-) create mode 100644 tests-unit/utils/extra_config_test.py create mode 100644 utils/__init__.py create mode 100644 utils/extra_config.py diff --git a/main.py b/main.py index d791a169c..a1db97cd0 100644 --- a/main.py +++ b/main.py @@ -63,6 +63,7 @@ import threading import gc import logging +from utils import extra_config if os.name == "nt": logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage()) @@ -85,7 +86,6 @@ if args.windows_standalone_build: pass import comfy.utils -import yaml import execution import server @@ -180,27 +180,6 @@ def cleanup_temp(): shutil.rmtree(temp_dir, ignore_errors=True) -def load_extra_path_config(yaml_path): - with open(yaml_path, 'r') as stream: - config = yaml.safe_load(stream) - for c in config: - conf = config[c] - if conf is None: - continue - base_path = None - if "base_path" in conf: - base_path = conf.pop("base_path") - for x in conf: - for y in conf[x].split("\n"): - if len(y) == 0: - continue - full_path = y - if base_path is not None: - full_path = os.path.join(base_path, full_path) - logging.info("Adding extra search path {} {}".format(x, full_path)) - folder_paths.add_model_folder_path(x, full_path) - - if __name__ == "__main__": if args.temp_directory: temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp") @@ -222,11 +201,11 @@ if __name__ == "__main__": extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml") if os.path.isfile(extra_model_paths_config_path): - load_extra_path_config(extra_model_paths_config_path) + extra_config.load_extra_path_config(extra_model_paths_config_path) if args.extra_model_paths_config: for config_path in itertools.chain(*args.extra_model_paths_config): - load_extra_path_config(config_path) + extra_config.load_extra_path_config(config_path) nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes) diff --git a/tests-unit/utils/extra_config_test.py b/tests-unit/utils/extra_config_test.py new file mode 100644 index 000000000..f56dd3e2e --- /dev/null +++ b/tests-unit/utils/extra_config_test.py @@ -0,0 +1,69 @@ +import pytest +import yaml +import os +from unittest.mock import Mock, patch, mock_open + +from utils.extra_config import load_extra_path_config +import folder_paths + +@pytest.fixture +def mock_yaml_content(): + return { + 'test_config': { + 'base_path': '~/App/', + 'checkpoints': 'subfolder1', + } + } + +@pytest.fixture +def mock_expanded_home(): + return '/home/user' + +@pytest.fixture +def mock_add_model_folder_path(): + return Mock() + +@pytest.fixture +def mock_expanduser(mock_expanded_home): + def _expanduser(path): + if path.startswith('~/'): + return os.path.join(mock_expanded_home, path[2:]) + return path + return _expanduser + +@pytest.fixture +def mock_yaml_safe_load(mock_yaml_content): + return Mock(return_value=mock_yaml_content) + +@patch('builtins.open', new_callable=mock_open, read_data="dummy file content") +def test_load_extra_model_paths_expands_userpath( + mock_file, + monkeypatch, + mock_add_model_folder_path, + mock_expanduser, + mock_yaml_safe_load, + mock_expanded_home +): + # Attach mocks used by load_extra_path_config + monkeypatch.setattr(folder_paths, 'add_model_folder_path', mock_add_model_folder_path) + monkeypatch.setattr(os.path, 'expanduser', mock_expanduser) + monkeypatch.setattr(yaml, 'safe_load', mock_yaml_safe_load) + + dummy_yaml_file_name = 'dummy_path.yaml' + load_extra_path_config(dummy_yaml_file_name) + + expected_calls = [ + ('checkpoints', os.path.join(mock_expanded_home, 'App', 'subfolder1')), + ] + + assert mock_add_model_folder_path.call_count == len(expected_calls) + + # Check if add_model_folder_path was called with the correct arguments + for actual_call, expected_call in zip(mock_add_model_folder_path.call_args_list, expected_calls): + assert actual_call.args == expected_call + + # Check if yaml.safe_load was called + mock_yaml_safe_load.assert_called_once() + + # Check if open was called with the correct file path + mock_file.assert_called_once_with(dummy_yaml_file_name, 'r') diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/utils/extra_config.py b/utils/extra_config.py new file mode 100644 index 000000000..23c2d791c --- /dev/null +++ b/utils/extra_config.py @@ -0,0 +1,25 @@ +import os +import yaml +import folder_paths +import logging + +def load_extra_path_config(yaml_path): + with open(yaml_path, 'r') as stream: + config = yaml.safe_load(stream) + for c in config: + conf = config[c] + if conf is None: + continue + base_path = None + if "base_path" in conf: + base_path = conf.pop("base_path") + base_path = os.path.expanduser(base_path) + for x in conf: + for y in conf[x].split("\n"): + if len(y) == 0: + continue + full_path = y + if base_path is not None: + full_path = os.path.join(base_path, full_path) + logging.info("Adding extra search path {} {}".format(x, full_path)) + folder_paths.add_model_folder_path(x, full_path) From bc94662b31351ac3ccde903c3829e5abce3d7b6f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 10 Sep 2024 00:43:37 -0400 Subject: [PATCH 04/26] Cleanup. --- main.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/main.py b/main.py index a1db97cd0..748846b4a 100644 --- a/main.py +++ b/main.py @@ -63,7 +63,7 @@ import threading import gc import logging -from utils import extra_config +import utils.extra_config if os.name == "nt": logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage()) @@ -201,11 +201,11 @@ if __name__ == "__main__": extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml") if os.path.isfile(extra_model_paths_config_path): - extra_config.load_extra_path_config(extra_model_paths_config_path) + utils.extra_config.load_extra_path_config(extra_model_paths_config_path) if args.extra_model_paths_config: for config_path in itertools.chain(*args.extra_model_paths_config): - extra_config.load_extra_path_config(config_path) + utils.extra_config.load_extra_path_config(config_path) nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes) From 81778a7feb233ccba046555d1cc90e7455df93ea Mon Sep 17 00:00:00 2001 From: Yoland Yan <4950057+yoland68@users.noreply.github.com> Date: Tue, 10 Sep 2024 13:44:49 +0900 Subject: [PATCH 05/26] =?UTF-8?q?[=F0=9F=97=BB=20Mount=20Fuji=20Commit]=20?= =?UTF-8?q?Add=20unit=20tests=20for=20folder=20path=20utilities=20(#4869)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All past 30 min of comtts are done on the top of Mt Fuji By Comfy, Robin, and Yoland All other comfy org members died on the way Introduced unit tests to verify the correctness of various folder path utility functions such as `get_directory_by_type`, `annotated_filepath`, and `recursive_search` among others. These tests cover scenarios including directory retrieval, filepath annotation, recursive file searches, and filtering files by extensions, enhancing the robustness and reliability of the codebase. --- tests-unit/comfy_test/folder_path_test.py | 66 +++++++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 tests-unit/comfy_test/folder_path_test.py diff --git a/tests-unit/comfy_test/folder_path_test.py b/tests-unit/comfy_test/folder_path_test.py new file mode 100644 index 000000000..0bbec593b --- /dev/null +++ b/tests-unit/comfy_test/folder_path_test.py @@ -0,0 +1,66 @@ +### 🗻 This file is created through the spirit of Mount Fuji at its peak +# TODO(yoland): clean up this after I get back down +import pytest +import os +import tempfile +from unittest.mock import patch + +import folder_paths + +@pytest.fixture +def temp_dir(): + with tempfile.TemporaryDirectory() as tmpdirname: + yield tmpdirname + + +def test_get_directory_by_type(): + test_dir = "/test/dir" + folder_paths.set_output_directory(test_dir) + assert folder_paths.get_directory_by_type("output") == test_dir + assert folder_paths.get_directory_by_type("invalid") is None + +def test_annotated_filepath(): + assert folder_paths.annotated_filepath("test.txt") == ("test.txt", None) + assert folder_paths.annotated_filepath("test.txt [output]") == ("test.txt", folder_paths.get_output_directory()) + assert folder_paths.annotated_filepath("test.txt [input]") == ("test.txt", folder_paths.get_input_directory()) + assert folder_paths.annotated_filepath("test.txt [temp]") == ("test.txt", folder_paths.get_temp_directory()) + +def test_get_annotated_filepath(): + default_dir = "/default/dir" + assert folder_paths.get_annotated_filepath("test.txt", default_dir) == os.path.join(default_dir, "test.txt") + assert folder_paths.get_annotated_filepath("test.txt [output]") == os.path.join(folder_paths.get_output_directory(), "test.txt") + +def test_add_model_folder_path(): + folder_paths.add_model_folder_path("test_folder", "/test/path") + assert "/test/path" in folder_paths.get_folder_paths("test_folder") + +def test_recursive_search(temp_dir): + os.makedirs(os.path.join(temp_dir, "subdir")) + open(os.path.join(temp_dir, "file1.txt"), "w").close() + open(os.path.join(temp_dir, "subdir", "file2.txt"), "w").close() + + files, dirs = folder_paths.recursive_search(temp_dir) + assert set(files) == {"file1.txt", os.path.join("subdir", "file2.txt")} + assert len(dirs) == 2 # temp_dir and subdir + +def test_filter_files_extensions(): + files = ["file1.txt", "file2.jpg", "file3.png", "file4.txt"] + assert folder_paths.filter_files_extensions(files, [".txt"]) == ["file1.txt", "file4.txt"] + assert folder_paths.filter_files_extensions(files, [".jpg", ".png"]) == ["file2.jpg", "file3.png"] + assert folder_paths.filter_files_extensions(files, []) == files + +@patch("folder_paths.recursive_search") +@patch("folder_paths.folder_names_and_paths") +def test_get_filename_list(mock_folder_names_and_paths, mock_recursive_search): + mock_folder_names_and_paths.__getitem__.return_value = (["/test/path"], {".txt"}) + mock_recursive_search.return_value = (["file1.txt", "file2.jpg"], {}) + assert folder_paths.get_filename_list("test_folder") == ["file1.txt"] + +def test_get_save_image_path(temp_dir): + with patch("folder_paths.output_directory", temp_dir): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path("test", temp_dir, 100, 100) + assert os.path.samefile(full_output_folder, temp_dir) + assert filename == "test" + assert counter == 1 + assert subfolder == "" + assert filename_prefix == "test" \ No newline at end of file From 36c83cdbba89a01830816240553f16dd72377cd9 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 11 Sep 2024 01:00:31 -0400 Subject: [PATCH 06/26] Limit origin check to when host is loopback. This should still prevent the exploit without breaking things for people who use reverse proxies. --- server.py | 34 +++++++++++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/server.py b/server.py index a1fca49cc..a2c078e40 100644 --- a/server.py +++ b/server.py @@ -12,6 +12,8 @@ import json import glob import struct import ssl +import socket +import ipaddress from PIL import Image, ImageOps from PIL.PngImagePlugin import PngInfo from io import BytesIO @@ -80,6 +82,32 @@ def create_cors_middleware(allowed_origin: str): return cors_middleware +def is_loopback(host): + if host is None: + return False + try: + if ipaddress.ip_address(host).is_loopback: + return True + else: + return False + except: + pass + + loopback = False + for family in (socket.AF_INET, socket.AF_INET6): + try: + r = socket.getaddrinfo(host, None, family, socket.SOCK_STREAM) + for family, _, _, _, sockaddr in r: + if not ipaddress.ip_address(sockaddr[0]).is_loopback: + return loopback + else: + loopback = True + except socket.gaierror: + pass + + return loopback + + def create_origin_only_middleware(): @web.middleware async def origin_only_middleware(request: web.Request, handler): @@ -93,12 +121,16 @@ def create_origin_only_middleware(): parsed = urllib.parse.urlparse(origin) origin_domain = parsed.netloc.lower() host_domain_parsed = urllib.parse.urlsplit('//' + host_domain) + + #limit the check to when the host domain is localhost, this makes it slightly less safe but should still prevent the exploit + loopback = is_loopback(host_domain_parsed.hostname) + if parsed.port is None: #if origin doesn't have a port strip it from the host to handle weird browsers, same for host host_domain = host_domain_parsed.hostname if host_domain_parsed.port is None: origin_domain = parsed.hostname - if host_domain is not None and origin_domain is not None and len(host_domain) > 0 and len(origin_domain) > 0: + if loopback and host_domain is not None and origin_domain is not None and len(host_domain) > 0 and len(origin_domain) > 0: if host_domain != origin_domain: logging.warning("WARNING: request with non matching host and origin {} != {}, returning 403".format(host_domain, origin_domain)) return web.Response(status=403) From e760bf5c4020753a5face9eaeb3ecede9a7b203d Mon Sep 17 00:00:00 2001 From: bymyself Date: Tue, 10 Sep 2024 23:00:07 -0700 Subject: [PATCH 07/26] Add content-type filter method to folder_paths (#4054) * Add content-type filter method to folder_paths * Add unit tests * Hardcode webp content-type * Annotate content_types as Literal["image", "video", "audio"] --- comfy_extras/nodes_audio.py | 9 +--- folder_paths.py | 28 ++++++++++ tests-unit/folder_paths_test/__init__.py | 0 .../filter_by_content_types_test.py | 52 +++++++++++++++++++ 4 files changed, 81 insertions(+), 8 deletions(-) create mode 100644 tests-unit/folder_paths_test/__init__.py create mode 100644 tests-unit/folder_paths_test/filter_by_content_types_test.py diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py index 762b48279..6990b3f9a 100644 --- a/comfy_extras/nodes_audio.py +++ b/comfy_extras/nodes_audio.py @@ -183,17 +183,10 @@ class PreviewAudio(SaveAudio): } class LoadAudio: - SUPPORTED_FORMATS = ('.wav', '.mp3', '.ogg', '.flac', '.aiff', '.aif') - @classmethod def INPUT_TYPES(s): input_dir = folder_paths.get_input_directory() - files = [ - f for f in os.listdir(input_dir) - if (os.path.isfile(os.path.join(input_dir, f)) - and f.endswith(LoadAudio.SUPPORTED_FORMATS) - ) - ] + files = folder_paths.filter_files_content_types(os.listdir(input_dir), ["audio", "video"]) return {"required": {"audio": (sorted(files), {"audio_upload": True})}} CATEGORY = "audio" diff --git a/folder_paths.py b/folder_paths.py index b154448fc..d7a7e0c3d 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -2,7 +2,9 @@ from __future__ import annotations import os import time +import mimetypes import logging +from typing import Set, List, Dict, Tuple, Literal from collections.abc import Collection supported_pt_extensions: set[str] = {'.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl', '.sft'} @@ -44,6 +46,10 @@ user_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "user filename_list_cache: dict[str, tuple[list[str], dict[str, float], float]] = {} +extension_mimetypes_cache = { + "webp" : "image", +} + def map_legacy(folder_name: str) -> str: legacy = {"unet": "diffusion_models"} return legacy.get(folder_name, folder_name) @@ -89,6 +95,28 @@ def get_directory_by_type(type_name: str) -> str | None: return get_input_directory() return None +def filter_files_content_types(files: List[str], content_types: Literal["image", "video", "audio"]) -> List[str]: + """ + Example: + files = os.listdir(folder_paths.get_input_directory()) + filter_files_content_types(files, ["image", "audio", "video"]) + """ + global extension_mimetypes_cache + result = [] + for file in files: + extension = file.split('.')[-1] + if extension not in extension_mimetypes_cache: + mime_type, _ = mimetypes.guess_type(file, strict=False) + if not mime_type: + continue + content_type = mime_type.split('/')[0] + extension_mimetypes_cache[extension] = content_type + else: + content_type = extension_mimetypes_cache[extension] + + if content_type in content_types: + result.append(file) + return result # determine base_dir rely on annotation if name is 'filename.ext [annotation]' format # otherwise use default_path as base_dir diff --git a/tests-unit/folder_paths_test/__init__.py b/tests-unit/folder_paths_test/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests-unit/folder_paths_test/filter_by_content_types_test.py b/tests-unit/folder_paths_test/filter_by_content_types_test.py new file mode 100644 index 000000000..5941bfa94 --- /dev/null +++ b/tests-unit/folder_paths_test/filter_by_content_types_test.py @@ -0,0 +1,52 @@ +import pytest +import os +import tempfile +from folder_paths import filter_files_content_types + +@pytest.fixture(scope="module") +def file_extensions(): + return { + 'image': ['bmp', 'cdr', 'gif', 'heif', 'ico', 'jpeg', 'jpg', 'pcx', 'png', 'pnm', 'ppm', 'psd', 'sgi', 'svg', 'tiff', 'webp', 'xbm', 'xcf', 'xpm'], + 'audio': ['aif', 'aifc', 'aiff', 'au', 'awb', 'flac', 'm4a', 'mp2', 'mp3', 'ogg', 'sd2', 'smp', 'snd', 'wav'], + 'video': ['avi', 'flv', 'm2v', 'm4v', 'mj2', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ogv', 'qt', 'webm', 'wmv'] + } + + +@pytest.fixture(scope="module") +def mock_dir(file_extensions): + with tempfile.TemporaryDirectory() as directory: + for content_type, extensions in file_extensions.items(): + for extension in extensions: + with open(f"{directory}/sample_{content_type}.{extension}", "w") as f: + f.write(f"Sample {content_type} file in {extension} format") + yield directory + + +def test_categorizes_all_correctly(mock_dir, file_extensions): + files = os.listdir(mock_dir) + for content_type, extensions in file_extensions.items(): + filtered_files = filter_files_content_types(files, [content_type]) + for extension in extensions: + assert f"sample_{content_type}.{extension}" in filtered_files + + +def test_categorizes_all_uniquely(mock_dir, file_extensions): + files = os.listdir(mock_dir) + for content_type, extensions in file_extensions.items(): + filtered_files = filter_files_content_types(files, [content_type]) + assert len(filtered_files) == len(extensions) + + +def test_handles_bad_extensions(): + files = ["file1.txt", "file2.py", "file3.example", "file4.pdf", "file5.ini", "file6.doc", "file7.md"] + assert filter_files_content_types(files, ["image", "audio", "video"]) == [] + + +def test_handles_no_extension(): + files = ["file1", "file2", "file3", "file4", "file5", "file6", "file7"] + assert filter_files_content_types(files, ["image", "audio", "video"]) == [] + + +def test_handles_no_files(): + files = [] + assert filter_files_content_types(files, ["image", "audio", "video"]) == [] \ No newline at end of file From b5d0f2a9086e7852b5d7422ed8d4b9ce2c93f47e Mon Sep 17 00:00:00 2001 From: yoinked Date: Tue, 10 Sep 2024 23:49:44 -0700 Subject: [PATCH 08/26] Add CFG++ to DPM++ 2S Ancestral (#3871) * Update sampling.py * Update samplers.py * my bad * "fix" the sampler * Update samplers.py * i named it wrong * minor sampling improvements mainly using a dynamic rho value (hey this sounds a lot like smea!!!) * revert rho change rho? r? its just 1/2 --- comfy/k_diffusion/sampling.py | 45 +++++++++++++++++++++++++++++++++++ comfy/samplers.py | 2 +- 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index 385a25164..f9f9170c9 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -1101,3 +1101,48 @@ def sample_euler_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=No if sigmas[i + 1] > 0: x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up return x +@torch.no_grad() +def sample_dpmpp_2s_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): + if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST): + return sample_dpmpp_2s_ancestral_RF(model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler) + + """Ancestral sampling with DPM-Solver++(2S) second-order steps.""" + extra_args = {} if extra_args is None else extra_args + noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler + + temp = [0] + def post_cfg_function(args): + temp[0] = args["uncond_denoised"] + return args["denoised"] + + model_options = extra_args.get("model_options", {}).copy() + extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True) + + s_in = x.new_ones([x.shape[0]]) + sigma_fn = lambda t: t.neg().exp() + t_fn = lambda sigma: sigma.log().neg() + + for i in trange(len(sigmas) - 1, disable=disable): + denoised = model(x, sigmas[i] * s_in, **extra_args) + sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta) + if callback is not None: + callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) + if sigma_down == 0: + # Euler method + d = to_d(x, sigmas[i], temp[0]) + dt = sigma_down - sigmas[i] + x = denoised + d * sigma_down + else: + # DPM-Solver++(2S) + t, t_next = t_fn(sigmas[i]), t_fn(sigma_down) + # r = torch.sinh(1 + (2 - eta) * (t_next - t) / (t - t_fn(sigma_up))) works only on non-cfgpp, weird + r = 1 / 2 + h = t_next - t + s = t + r * h + x_2 = (sigma_fn(s) / sigma_fn(t)) * (x + (denoised - temp[0])) - (-h * r).expm1() * denoised + denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args) + x = (sigma_fn(t_next) / sigma_fn(t)) * (x + (denoised - temp[0])) - (-h).expm1() * denoised_2 + # Noise addition + if sigmas[i + 1] > 0: + x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up + return x diff --git a/comfy/samplers.py b/comfy/samplers.py index b278dae24..a20f65d4a 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -570,7 +570,7 @@ class Sampler: return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma KSAMPLER_NAMES = ["euler", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_cfg_pp", "heun", "heunpp2","dpm_2", "dpm_2_ancestral", - "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu", + "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_2s_ancestral_cfg_pp", "dpmpp_sde", "dpmpp_sde_gpu", "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "ipndm", "ipndm_v", "deis"] From 9f4daca9d9994c600ef713856e0182d0177b9a1e Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 11 Sep 2024 02:51:36 -0400 Subject: [PATCH 09/26] Doesn't really make sense for cfg_pp sampler to call regular one. --- comfy/k_diffusion/sampling.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index f9f9170c9..ae56aee8b 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -1103,13 +1103,10 @@ def sample_euler_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=No return x @torch.no_grad() def sample_dpmpp_2s_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None): - if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST): - return sample_dpmpp_2s_ancestral_RF(model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler) - """Ancestral sampling with DPM-Solver++(2S) second-order steps.""" extra_args = {} if extra_args is None else extra_args noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler - + temp = [0] def post_cfg_function(args): temp[0] = args["uncond_denoised"] From d247bc5a9cd624abe38fee7ac80ec03fda764845 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Thu, 12 Sep 2024 14:52:06 +0900 Subject: [PATCH 10/26] Expand variables in base_path for extra_config_paths.yaml. (#4893) * Expand variables in base_path for extra_config_paths.yaml. * Fix comments. --- tests-unit/utils/extra_config_test.py | 55 +++++++++++++++++++++++++++ utils/extra_config.py | 2 +- 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/tests-unit/utils/extra_config_test.py b/tests-unit/utils/extra_config_test.py index f56dd3e2e..06349560d 100644 --- a/tests-unit/utils/extra_config_test.py +++ b/tests-unit/utils/extra_config_test.py @@ -19,6 +19,24 @@ def mock_yaml_content(): def mock_expanded_home(): return '/home/user' +@pytest.fixture +def yaml_config_with_appdata(): + return """ + test_config: + base_path: '%APPDATA%/ComfyUI' + checkpoints: 'models/checkpoints' + """ + +@pytest.fixture +def mock_yaml_content_appdata(yaml_config_with_appdata): + return yaml.safe_load(yaml_config_with_appdata) + +@pytest.fixture +def mock_expandvars_appdata(): + mock = Mock() + mock.side_effect = lambda path: path.replace('%APPDATA%', 'C:/Users/TestUser/AppData/Roaming') + return mock + @pytest.fixture def mock_add_model_folder_path(): return Mock() @@ -67,3 +85,40 @@ def test_load_extra_model_paths_expands_userpath( # Check if open was called with the correct file path mock_file.assert_called_once_with(dummy_yaml_file_name, 'r') + +@patch('builtins.open', new_callable=mock_open) +def test_load_extra_model_paths_expands_appdata( + mock_file, + monkeypatch, + mock_add_model_folder_path, + mock_expandvars_appdata, + yaml_config_with_appdata, + mock_yaml_content_appdata +): + # Set the mock_file to return yaml with appdata as a variable + mock_file.return_value.read.return_value = yaml_config_with_appdata + + # Attach mocks + monkeypatch.setattr(folder_paths, 'add_model_folder_path', mock_add_model_folder_path) + monkeypatch.setattr(os.path, 'expandvars', mock_expandvars_appdata) + monkeypatch.setattr(yaml, 'safe_load', Mock(return_value=mock_yaml_content_appdata)) + + # Mock expanduser to do nothing (since we're not testing it here) + monkeypatch.setattr(os.path, 'expanduser', lambda x: x) + + dummy_yaml_file_name = 'dummy_path.yaml' + load_extra_path_config(dummy_yaml_file_name) + + expected_base_path = 'C:/Users/TestUser/AppData/Roaming/ComfyUI' + expected_calls = [ + ('checkpoints', os.path.join(expected_base_path, 'models/checkpoints')), + ] + + assert mock_add_model_folder_path.call_count == len(expected_calls) + + # Check the base path variable was expanded + for actual_call, expected_call in zip(mock_add_model_folder_path.call_args_list, expected_calls): + assert actual_call.args == expected_call + + # Verify that expandvars was called + assert mock_expandvars_appdata.called diff --git a/utils/extra_config.py b/utils/extra_config.py index 23c2d791c..7ac3650ac 100644 --- a/utils/extra_config.py +++ b/utils/extra_config.py @@ -13,7 +13,7 @@ def load_extra_path_config(yaml_path): base_path = None if "base_path" in conf: base_path = conf.pop("base_path") - base_path = os.path.expanduser(base_path) + base_path = os.path.expandvars(os.path.expanduser(base_path)) for x in conf: for y in conf[x].split("\n"): if len(y) == 0: From 9d720187f181bc5a74b7bd36b8ad6276b1336000 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 12 Sep 2024 03:56:52 -0400 Subject: [PATCH 11/26] types -> comfy_types to fix import issue. --- comfy/{types.py => comfy_types.py} | 0 comfy/model_patcher.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename comfy/{types.py => comfy_types.py} (100%) diff --git a/comfy/types.py b/comfy/comfy_types.py similarity index 100% rename from comfy/types.py rename to comfy/comfy_types.py diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 3f5d90273..a89ad4fd7 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -28,7 +28,7 @@ import comfy.utils import comfy.float import comfy.model_management import comfy.lora -from comfy.types import UnetWrapperFunction +from comfy.comfy_types import UnetWrapperFunction def string_to_seed(data): crc = 0xFFFFFFFF From 405b529545f56658960a28e0c73e2a9c64ce2376 Mon Sep 17 00:00:00 2001 From: Yoland Yan <4950057+yoland68@users.noreply.github.com> Date: Thu, 12 Sep 2024 17:53:08 +0900 Subject: [PATCH 12/26] Minor: update tests-unit README.md (#4896) --- tests-unit/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests-unit/README.md b/tests-unit/README.md index 94abd9853..81692b8f1 100644 --- a/tests-unit/README.md +++ b/tests-unit/README.md @@ -2,7 +2,7 @@ ## Install test dependencies -`pip install -r tests-units/requirements.txt` +`pip install -r tests-unit/requirements.txt` ## Run tests -`pytest tests-units/` +`pytest tests-unit/` From d0b7ab88ba0f1cb4ab16e0425f5229e60c934536 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 12 Sep 2024 05:23:32 -0400 Subject: [PATCH 13/26] Add a simple experimental TorchCompileModel node. It probably only works on Linux. For maximum speed on Flux with Nvidia 40 series/ada and newer try using this node with fp8_e4m3fn and the --fast argument. --- comfy_extras/nodes_torch_compile.py | 21 +++++++++++++++++++++ nodes.py | 1 + 2 files changed, 22 insertions(+) create mode 100644 comfy_extras/nodes_torch_compile.py diff --git a/comfy_extras/nodes_torch_compile.py b/comfy_extras/nodes_torch_compile.py new file mode 100644 index 000000000..1d914fa93 --- /dev/null +++ b/comfy_extras/nodes_torch_compile.py @@ -0,0 +1,21 @@ +import torch + +class TorchCompileModel: + @classmethod + def INPUT_TYPES(s): + return {"required": { "model": ("MODEL",), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + + CATEGORY = "_for_testing" + EXPERIMENTAL = True + + def patch(self, model): + m = model.clone() + m.add_object_patch("diffusion_model", torch.compile(model=m.get_model_object("diffusion_model"))) + return (m, ) + +NODE_CLASS_MAPPINGS = { + "TorchCompileModel": TorchCompileModel, +} diff --git a/nodes.py b/nodes.py index bbe73282a..1f14aaf11 100644 --- a/nodes.py +++ b/nodes.py @@ -2102,6 +2102,7 @@ def init_builtin_extra_nodes(): "nodes_hunyuan.py", "nodes_flux.py", "nodes_lora_extract.py", + "nodes_torch_compile.py", ] import_failed = [] From b962db9952d45b0ed8185ce7d18cf7d960f99de4 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Thu, 12 Sep 2024 21:10:27 +0900 Subject: [PATCH 14/26] Add cli arg to override user directory (#4856) * Override user directory. * Use overridden user directory. * Remove prints. * Remove references to global user_files. * Remove unused replace_folder function. * Remove newline. * Remove global during get_user_directory. * Add validation. --- app/user_manager.py | 17 +++++++++-------- comfy/cli_args.py | 2 ++ folder_paths.py | 7 +++++++ main.py | 5 +++++ 4 files changed, 23 insertions(+), 8 deletions(-) diff --git a/app/user_manager.py b/app/user_manager.py index 95ad157c3..62c22cde5 100644 --- a/app/user_manager.py +++ b/app/user_manager.py @@ -7,16 +7,15 @@ import shutil from aiohttp import web from urllib import parse from comfy.cli_args import args -from folder_paths import user_directory +import folder_paths from .app_settings import AppSettings default_user = "default" -users_file = os.path.join(user_directory, "users.json") class UserManager(): def __init__(self): - global user_directory + user_directory = folder_paths.get_user_directory() self.settings = AppSettings(self) if not os.path.exists(user_directory): @@ -26,14 +25,17 @@ class UserManager(): print("****** For multi-user setups add the --multi-user CLI argument to enable multiple user profiles. ******") if args.multi_user: - if os.path.isfile(users_file): - with open(users_file) as f: + if os.path.isfile(self.get_users_file()): + with open(self.get_users_file()) as f: self.users = json.load(f) else: self.users = {} else: self.users = {"default": "default"} + def get_users_file(self): + return os.path.join(folder_paths.get_user_directory(), "users.json") + def get_request_user_id(self, request): user = "default" if args.multi_user and "comfy-user" in request.headers: @@ -45,7 +47,7 @@ class UserManager(): return user def get_request_user_filepath(self, request, file, type="userdata", create_dir=True): - global user_directory + user_directory = folder_paths.get_user_directory() if type == "userdata": root_dir = user_directory @@ -85,8 +87,7 @@ class UserManager(): self.users[user_id] = name - global users_file - with open(users_file, "w") as f: + with open(self.get_users_file(), "w") as f: json.dump(self.users, f) return user_id diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 65b879907..d0f706432 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -171,6 +171,8 @@ parser.add_argument( help="The local filesystem path to the directory where the frontend is located. Overrides --front-end-version.", ) +parser.add_argument("--user-directory", type=is_valid_directory, default=None, help="Set the ComfyUI user directory with an absolute path.") + if comfy.options.args_parsing: args = parser.parse_args() else: diff --git a/folder_paths.py b/folder_paths.py index d7a7e0c3d..263704a4c 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -84,6 +84,13 @@ def get_input_directory() -> str: global input_directory return input_directory +def get_user_directory() -> str: + return user_directory + +def set_user_directory(user_dir: str) -> None: + global user_directory + user_directory = user_dir + #NOTE: used in http server so don't put folders that should not be accessed remotely def get_directory_by_type(type_name: str) -> str | None: diff --git a/main.py b/main.py index 748846b4a..00b49fc4c 100644 --- a/main.py +++ b/main.py @@ -232,6 +232,11 @@ if __name__ == "__main__": input_dir = os.path.abspath(args.input_directory) logging.info(f"Setting input directory to: {input_dir}") folder_paths.set_input_directory(input_dir) + + if args.user_directory: + user_dir = os.path.abspath(args.user_directory) + logging.info(f"Setting user directory to: {user_dir}") + folder_paths.set_user_directory(user_dir) if args.quick_test_for_ci: exit(0) From 7c6eb4fb29c2c3b4f12605003ab1b7c612d4b4e1 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 12 Sep 2024 20:27:07 -0400 Subject: [PATCH 15/26] Set some nodes as DEPRECATED. --- comfy_extras/nodes_perpneg.py | 1 + nodes.py | 1 + 2 files changed, 2 insertions(+) diff --git a/comfy_extras/nodes_perpneg.py b/comfy_extras/nodes_perpneg.py index 546276aa1..762c40220 100644 --- a/comfy_extras/nodes_perpneg.py +++ b/comfy_extras/nodes_perpneg.py @@ -26,6 +26,7 @@ class PerpNeg: FUNCTION = "patch" CATEGORY = "_for_testing" + DEPRECATED = True def patch(self, model, empty_conditioning, neg_scale): m = model.clone() diff --git a/nodes.py b/nodes.py index 1f14aaf11..0d3749d42 100644 --- a/nodes.py +++ b/nodes.py @@ -511,6 +511,7 @@ class CheckpointLoader: FUNCTION = "load_checkpoint" CATEGORY = "advanced/loaders" + DEPRECATED = True def load_checkpoint(self, config_name, ckpt_name): config_path = folder_paths.get_full_path("configs", config_name) From f6b7194f64ab4b32018eea0da9d9d89a30b582aa Mon Sep 17 00:00:00 2001 From: JettHu <35261585+JettHu@users.noreply.github.com> Date: Fri, 13 Sep 2024 11:02:52 +0800 Subject: [PATCH 16/26] Reduce repeated calls of get_immediate_node_signature for ancestors in cache (#4871) --- comfy_execution/caching.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index 62311ed7f..f163bc989 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -67,6 +67,7 @@ class CacheKeySetInputSignature(CacheKeySet): super().__init__(dynprompt, node_ids, is_changed_cache) self.dynprompt = dynprompt self.is_changed_cache = is_changed_cache + self.immediate_node_signature = {} self.add_keys(node_ids) def include_node_id_in_input(self) -> bool: @@ -94,6 +95,8 @@ class CacheKeySetInputSignature(CacheKeySet): if not dynprompt.has_node(node_id): # This node doesn't exist -- we can't cache it. return [float("NaN")] + if node_id in self.immediate_node_signature: # reduce repeated calls of ancestors + return self.immediate_node_signature[node_id] node = dynprompt.get_node(node_id) class_type = node["class_type"] class_def = nodes.NODE_CLASS_MAPPINGS[class_type] @@ -108,6 +111,7 @@ class CacheKeySetInputSignature(CacheKeySet): signature.append((key,("ANCESTOR", ancestor_index, ancestor_socket))) else: signature.append((key, inputs[key])) + self.immediate_node_signature[node_id] = signature return signature # This function returns a list of all ancestors of the given node. The order of the list is From cb12ad7049e98e3621c6326d9f139c8330799bd2 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Fri, 13 Sep 2024 15:40:59 +0900 Subject: [PATCH 17/26] Add full_info flag in /userdata endpoint to list out file size and last modified timestamp (#4905) * Add full_info flag in /userdata endpoint to list out file size and last modified timestamp * nit --- app/user_manager.py | 33 +++++-- .../prompt_server_test/user_manager_test.py | 90 +++++++++++++++++++ 2 files changed, 115 insertions(+), 8 deletions(-) create mode 100644 tests-unit/prompt_server_test/user_manager_test.py diff --git a/app/user_manager.py b/app/user_manager.py index 62c22cde5..260c383b4 100644 --- a/app/user_manager.py +++ b/app/user_manager.py @@ -121,21 +121,38 @@ class UserManager(): directory = request.rel_url.query.get('dir', '') if not directory: return web.Response(status=400) - + path = self.get_request_user_filepath(request, directory) if not path: return web.Response(status=403) - + if not os.path.exists(path): return web.Response(status=404) - + recurse = request.rel_url.query.get('recurse', '').lower() == "true" - results = glob.glob(os.path.join( - glob.escape(path), '**/*'), recursive=recurse) - results = [os.path.relpath(x, path) for x in results if os.path.isfile(x)] - + full_info = request.rel_url.query.get('full_info', '').lower() == "true" + + # Use different patterns based on whether we're recursing or not + if recurse: + pattern = os.path.join(glob.escape(path), '**', '*') + else: + pattern = os.path.join(glob.escape(path), '*') + + results = glob.glob(pattern, recursive=recurse) + + if full_info: + results = [ + { + 'path': os.path.relpath(x, path), + 'size': os.path.getsize(x), + 'modified': os.path.getmtime(x) + } for x in results if os.path.isfile(x) + ] + else: + results = [os.path.relpath(x, path) for x in results if os.path.isfile(x)] + split_path = request.rel_url.query.get('split', '').lower() == "true" - if split_path: + if split_path and not full_info: results = [[x] + x.split(os.sep) for x in results] return web.json_response(results) diff --git a/tests-unit/prompt_server_test/user_manager_test.py b/tests-unit/prompt_server_test/user_manager_test.py new file mode 100644 index 000000000..c71050a2f --- /dev/null +++ b/tests-unit/prompt_server_test/user_manager_test.py @@ -0,0 +1,90 @@ +import pytest +import os +from aiohttp import web +from app.user_manager import UserManager + +pytestmark = ( + pytest.mark.asyncio +) # This applies the asyncio mark to all test functions in the module + + +@pytest.fixture +def user_manager(tmp_path): + um = UserManager() + um.get_request_user_filepath = lambda req, file, **kwargs: os.path.join( + tmp_path, file + ) + return um + + +@pytest.fixture +def app(user_manager): + app = web.Application() + routes = web.RouteTableDef() + user_manager.add_routes(routes) + app.add_routes(routes) + return app + + +async def test_listuserdata_empty_directory(aiohttp_client, app, tmp_path): + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir") + assert resp.status == 404 + + +async def test_listuserdata_with_files(aiohttp_client, app, tmp_path): + os.makedirs(tmp_path / "test_dir") + with open(tmp_path / "test_dir" / "file1.txt", "w") as f: + f.write("test content") + + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir") + assert resp.status == 200 + assert await resp.json() == ["file1.txt"] + + +async def test_listuserdata_recursive(aiohttp_client, app, tmp_path): + os.makedirs(tmp_path / "test_dir" / "subdir") + with open(tmp_path / "test_dir" / "file1.txt", "w") as f: + f.write("test content") + with open(tmp_path / "test_dir" / "subdir" / "file2.txt", "w") as f: + f.write("test content") + + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir&recurse=true") + assert resp.status == 200 + assert set(await resp.json()) == {"file1.txt", os.path.join("subdir", "file2.txt")} + + +async def test_listuserdata_full_info(aiohttp_client, app, tmp_path): + os.makedirs(tmp_path / "test_dir") + with open(tmp_path / "test_dir" / "file1.txt", "w") as f: + f.write("test content") + + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir&full_info=true") + assert resp.status == 200 + result = await resp.json() + assert len(result) == 1 + assert result[0]["path"] == "file1.txt" + assert "size" in result[0] + assert "modified" in result[0] + + +async def test_listuserdata_split_path(aiohttp_client, app, tmp_path): + os.makedirs(tmp_path / "test_dir" / "subdir") + with open(tmp_path / "test_dir" / "subdir" / "file1.txt", "w") as f: + f.write("test content") + + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir&recurse=true&split=true") + assert resp.status == 200 + assert await resp.json() == [ + [os.path.join("subdir", "file1.txt"), "subdir", "file1.txt"] + ] + + +async def test_listuserdata_invalid_directory(aiohttp_client, app): + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=") + assert resp.status == 400 From d2247c1e6130a940bf702f30289fdf71d00b53b3 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Fri, 13 Sep 2024 16:45:31 +0900 Subject: [PATCH 18/26] Normalize path returned by /userdata to always use / as separator (#4906) --- app/user_manager.py | 16 +++++---- .../prompt_server_test/user_manager_test.py | 34 +++++++++++++++++-- 2 files changed, 42 insertions(+), 8 deletions(-) diff --git a/app/user_manager.py b/app/user_manager.py index 260c383b4..42bc496d5 100644 --- a/app/user_manager.py +++ b/app/user_manager.py @@ -120,14 +120,14 @@ class UserManager(): async def listuserdata(request): directory = request.rel_url.query.get('dir', '') if not directory: - return web.Response(status=400) + return web.Response(status=400, text="Directory not provided") path = self.get_request_user_filepath(request, directory) if not path: - return web.Response(status=403) + return web.Response(status=403, text="Invalid directory") if not os.path.exists(path): - return web.Response(status=404) + return web.Response(status=404, text="Directory not found") recurse = request.rel_url.query.get('recurse', '').lower() == "true" full_info = request.rel_url.query.get('full_info', '').lower() == "true" @@ -143,17 +143,21 @@ class UserManager(): if full_info: results = [ { - 'path': os.path.relpath(x, path), + 'path': os.path.relpath(x, path).replace(os.sep, '/'), 'size': os.path.getsize(x), 'modified': os.path.getmtime(x) } for x in results if os.path.isfile(x) ] else: - results = [os.path.relpath(x, path) for x in results if os.path.isfile(x)] + results = [ + os.path.relpath(x, path).replace(os.sep, '/') + for x in results + if os.path.isfile(x) + ] split_path = request.rel_url.query.get('split', '').lower() == "true" if split_path and not full_info: - results = [[x] + x.split(os.sep) for x in results] + results = [[x] + x.split('/') for x in results] return web.json_response(results) diff --git a/tests-unit/prompt_server_test/user_manager_test.py b/tests-unit/prompt_server_test/user_manager_test.py index c71050a2f..936c6bd27 100644 --- a/tests-unit/prompt_server_test/user_manager_test.py +++ b/tests-unit/prompt_server_test/user_manager_test.py @@ -2,6 +2,7 @@ import pytest import os from aiohttp import web from app.user_manager import UserManager +from unittest.mock import patch pytestmark = ( pytest.mark.asyncio @@ -53,7 +54,7 @@ async def test_listuserdata_recursive(aiohttp_client, app, tmp_path): client = await aiohttp_client(app) resp = await client.get("/userdata?dir=test_dir&recurse=true") assert resp.status == 200 - assert set(await resp.json()) == {"file1.txt", os.path.join("subdir", "file2.txt")} + assert set(await resp.json()) == {"file1.txt", "subdir/file2.txt"} async def test_listuserdata_full_info(aiohttp_client, app, tmp_path): @@ -80,7 +81,7 @@ async def test_listuserdata_split_path(aiohttp_client, app, tmp_path): resp = await client.get("/userdata?dir=test_dir&recurse=true&split=true") assert resp.status == 200 assert await resp.json() == [ - [os.path.join("subdir", "file1.txt"), "subdir", "file1.txt"] + ["subdir/file1.txt", "subdir", "file1.txt"] ] @@ -88,3 +89,32 @@ async def test_listuserdata_invalid_directory(aiohttp_client, app): client = await aiohttp_client(app) resp = await client.get("/userdata?dir=") assert resp.status == 400 + + +async def test_listuserdata_normalized_separator(aiohttp_client, app, tmp_path): + os_sep = "\\" + with patch("os.sep", os_sep): + with patch("os.path.sep", os_sep): + os.makedirs(tmp_path / "test_dir" / "subdir") + with open(tmp_path / "test_dir" / "subdir" / "file1.txt", "w") as f: + f.write("test content") + + client = await aiohttp_client(app) + resp = await client.get("/userdata?dir=test_dir&recurse=true") + assert resp.status == 200 + result = await resp.json() + assert len(result) == 1 + assert "/" in result[0] # Ensure forward slash is used + assert "\\" not in result[0] # Ensure backslash is not present + assert result[0] == "subdir/file1.txt" + + # Test with full_info + resp = await client.get( + "/userdata?dir=test_dir&recurse=true&full_info=true" + ) + assert resp.status == 200 + result = await resp.json() + assert len(result) == 1 + assert "/" in result[0]["path"] # Ensure forward slash is used + assert "\\" not in result[0]["path"] # Ensure backslash is not present + assert result[0]["path"] == "subdir/file1.txt" From 6fb44c4b7cdf88e91e8ec53c828d1ca5e3fe2b0d Mon Sep 17 00:00:00 2001 From: Acly Date: Fri, 13 Sep 2024 14:25:11 +0200 Subject: [PATCH 19/26] Make adding links/nodes to ExecutionList non-recursive (#4886) Graphs with 300+ chained nodes run into maximum recursion depth error (limit is 1000 in CPython) --- comfy_execution/graph.py | 65 +++++++++++++++++++++++----------------- 1 file changed, 38 insertions(+), 27 deletions(-) diff --git a/comfy_execution/graph.py b/comfy_execution/graph.py index b53e10f3f..0b5bf1899 100644 --- a/comfy_execution/graph.py +++ b/comfy_execution/graph.py @@ -99,30 +99,44 @@ class TopologicalSort: self.add_strong_link(from_node_id, from_socket, to_node_id) def add_strong_link(self, from_node_id, from_socket, to_node_id): - self.add_node(from_node_id) - if to_node_id not in self.blocking[from_node_id]: - self.blocking[from_node_id][to_node_id] = {} - self.blockCount[to_node_id] += 1 - self.blocking[from_node_id][to_node_id][from_socket] = True + if not self.is_cached(from_node_id): + self.add_node(from_node_id) + if to_node_id not in self.blocking[from_node_id]: + self.blocking[from_node_id][to_node_id] = {} + self.blockCount[to_node_id] += 1 + self.blocking[from_node_id][to_node_id][from_socket] = True - def add_node(self, unique_id, include_lazy=False, subgraph_nodes=None): - if unique_id in self.pendingNodes: - return - self.pendingNodes[unique_id] = True - self.blockCount[unique_id] = 0 - self.blocking[unique_id] = {} + def add_node(self, node_unique_id, include_lazy=False, subgraph_nodes=None): + node_ids = [node_unique_id] + links = [] - inputs = self.dynprompt.get_node(unique_id)["inputs"] - for input_name in inputs: - value = inputs[input_name] - if is_link(value): - from_node_id, from_socket = value - if subgraph_nodes is not None and from_node_id not in subgraph_nodes: - continue - input_type, input_category, input_info = self.get_input_info(unique_id, input_name) - is_lazy = input_info is not None and "lazy" in input_info and input_info["lazy"] - if include_lazy or not is_lazy: - self.add_strong_link(from_node_id, from_socket, unique_id) + while len(node_ids) > 0: + unique_id = node_ids.pop() + if unique_id in self.pendingNodes: + continue + + self.pendingNodes[unique_id] = True + self.blockCount[unique_id] = 0 + self.blocking[unique_id] = {} + + inputs = self.dynprompt.get_node(unique_id)["inputs"] + for input_name in inputs: + value = inputs[input_name] + if is_link(value): + from_node_id, from_socket = value + if subgraph_nodes is not None and from_node_id not in subgraph_nodes: + continue + input_type, input_category, input_info = self.get_input_info(unique_id, input_name) + is_lazy = input_info is not None and "lazy" in input_info and input_info["lazy"] + if (include_lazy or not is_lazy) and not self.is_cached(from_node_id): + node_ids.append(from_node_id) + links.append((from_node_id, from_socket, unique_id)) + + for link in links: + self.add_strong_link(*link) + + def is_cached(self, node_id): + return False def get_ready_nodes(self): return [node_id for node_id in self.pendingNodes if self.blockCount[node_id] == 0] @@ -146,11 +160,8 @@ class ExecutionList(TopologicalSort): self.output_cache = output_cache self.staged_node_id = None - def add_strong_link(self, from_node_id, from_socket, to_node_id): - if self.output_cache.get(from_node_id) is not None: - # Nothing to do - return - super().add_strong_link(from_node_id, from_socket, to_node_id) + def is_cached(self, node_id): + return self.output_cache.get(node_id) is not None def stage_node_execution(self): assert self.staged_node_id is None From cf80d28689c1dfa8d22d7fa8e68db7ea3de0fd37 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 13 Sep 2024 09:52:20 -0400 Subject: [PATCH 20/26] Support loading controlnets with different input. --- comfy/controlnet.py | 4 +++- comfy/ldm/flux/controlnet.py | 9 +++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/comfy/controlnet.py b/comfy/controlnet.py index c0f9b6511..860891965 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -449,7 +449,9 @@ def load_controlnet_flux_instantx(sd): if union_cnet in new_sd: num_union_modes = new_sd[union_cnet].shape[0] - control_model = comfy.ldm.flux.controlnet.ControlNetFlux(latent_input=True, num_union_modes=num_union_modes, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) + control_latent_channels = new_sd.get("pos_embed_input.weight").shape[1] // 4 + + control_model = comfy.ldm.flux.controlnet.ControlNetFlux(latent_input=True, num_union_modes=num_union_modes, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) control_model = controlnet_load_state_dict(control_model, new_sd) latent_format = comfy.latent_formats.Flux() diff --git a/comfy/ldm/flux/controlnet.py b/comfy/ldm/flux/controlnet.py index d8b776129..c033dea52 100644 --- a/comfy/ldm/flux/controlnet.py +++ b/comfy/ldm/flux/controlnet.py @@ -52,7 +52,7 @@ class MistolineControlnetBlock(nn.Module): class ControlNetFlux(Flux): - def __init__(self, latent_input=False, num_union_modes=0, mistoline=False, image_model=None, dtype=None, device=None, operations=None, **kwargs): + def __init__(self, latent_input=False, num_union_modes=0, mistoline=False, control_latent_channels=None, image_model=None, dtype=None, device=None, operations=None, **kwargs): super().__init__(final_layer=False, dtype=dtype, device=device, operations=operations, **kwargs) self.main_model_double = 19 @@ -80,7 +80,12 @@ class ControlNetFlux(Flux): self.gradient_checkpointing = False self.latent_input = latent_input - self.pos_embed_input = operations.Linear(self.in_channels, self.hidden_size, bias=True, dtype=dtype, device=device) + if control_latent_channels is None: + control_latent_channels = self.in_channels + else: + control_latent_channels *= 2 * 2 #patch size + + self.pos_embed_input = operations.Linear(control_latent_channels, self.hidden_size, bias=True, dtype=dtype, device=device) if not self.latent_input: if self.mistoline: self.input_cond_block = MistolineCondDownsamplBlock(dtype=dtype, device=device, operations=operations) From b3ce8fb9fd6209ef7542c90daac7f1cc8a969b70 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 13 Sep 2024 23:24:47 -0400 Subject: [PATCH 21/26] Revert "Reduce repeated calls of get_immediate_node_signature for ancestors in cache (#4871)" This reverts commit f6b7194f64ab4b32018eea0da9d9d89a30b582aa. --- comfy_execution/caching.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index f163bc989..62311ed7f 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -67,7 +67,6 @@ class CacheKeySetInputSignature(CacheKeySet): super().__init__(dynprompt, node_ids, is_changed_cache) self.dynprompt = dynprompt self.is_changed_cache = is_changed_cache - self.immediate_node_signature = {} self.add_keys(node_ids) def include_node_id_in_input(self) -> bool: @@ -95,8 +94,6 @@ class CacheKeySetInputSignature(CacheKeySet): if not dynprompt.has_node(node_id): # This node doesn't exist -- we can't cache it. return [float("NaN")] - if node_id in self.immediate_node_signature: # reduce repeated calls of ancestors - return self.immediate_node_signature[node_id] node = dynprompt.get_node(node_id) class_type = node["class_type"] class_def = nodes.NODE_CLASS_MAPPINGS[class_type] @@ -111,7 +108,6 @@ class CacheKeySetInputSignature(CacheKeySet): signature.append((key,("ANCESTOR", ancestor_index, ancestor_socket))) else: signature.append((key, inputs[key])) - self.immediate_node_signature[node_id] = signature return signature # This function returns a list of all ancestors of the given node. The order of the list is From 369a6dd2c499a02b42496bab6286e34e9350c740 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Sat, 14 Sep 2024 12:30:44 +0900 Subject: [PATCH 22/26] Remove empty spaces in user_manager.py (#4917) --- app/user_manager.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/app/user_manager.py b/app/user_manager.py index 42bc496d5..05e6805f7 100644 --- a/app/user_manager.py +++ b/app/user_manager.py @@ -7,7 +7,7 @@ import shutil from aiohttp import web from urllib import parse from comfy.cli_args import args -import folder_paths +import folder_paths from .app_settings import AppSettings default_user = "default" @@ -65,7 +65,7 @@ class UserManager(): # Check if filename is url encoded if "%" in file: file = parse.unquote(file) - + # prevent leaving /{type}/{user} path = os.path.abspath(os.path.join(user_root, file)) if os.path.commonpath((user_root, path)) != user_root: @@ -165,14 +165,14 @@ class UserManager(): file = request.match_info.get(param, None) if not file: return web.Response(status=400) - + path = self.get_request_user_filepath(request, file) if not path: return web.Response(status=403) - + if check_exists and not os.path.exists(path): return web.Response(status=404) - + return path @routes.get("/userdata/{file}") @@ -180,7 +180,7 @@ class UserManager(): path = get_user_data_path(request, check_exists=True) if not isinstance(path, str): return path - + return web.FileResponse(path) @routes.post("/userdata/{file}") @@ -188,7 +188,7 @@ class UserManager(): path = get_user_data_path(request) if not isinstance(path, str): return path - + overwrite = request.query["overwrite"] != "false" if not overwrite and os.path.exists(path): return web.Response(status=409) @@ -197,7 +197,7 @@ class UserManager(): with open(path, "wb") as f: f.write(body) - + resp = os.path.relpath(path, self.get_request_user_filepath(request, None)) return web.json_response(resp) @@ -208,7 +208,7 @@ class UserManager(): return path os.remove(path) - + return web.Response(status=204) @routes.post("/userdata/{file}/move/{dest}") @@ -216,17 +216,17 @@ class UserManager(): source = get_user_data_path(request, check_exists=True) if not isinstance(source, str): return source - + dest = get_user_data_path(request, check_exists=False, param="dest") if not isinstance(source, str): return dest - + overwrite = request.query["overwrite"] != "false" if not overwrite and os.path.exists(dest): return web.Response(status=409) print(f"moving '{source}' -> '{dest}'") shutil.move(source, dest) - + resp = os.path.relpath(dest, self.get_request_user_filepath(request, None)) return web.json_response(resp) From f48e390032f8d27a450a35ef4aa4b775b078cbf9 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 14 Sep 2024 09:05:16 -0400 Subject: [PATCH 23/26] Support AliMama SD3 and Flux inpaint controlnets. Use the ControlNetInpaintingAliMamaApply node. --- comfy/cldm/mmdit.py | 6 ++++- comfy/controlnet.py | 40 +++++++++++++++++++++++++------- comfy_extras/nodes_controlnet.py | 29 +++++++++++++++++++++++ nodes.py | 4 ++-- 4 files changed, 68 insertions(+), 11 deletions(-) diff --git a/comfy/cldm/mmdit.py b/comfy/cldm/mmdit.py index 025c2fb5d..54a58ab83 100644 --- a/comfy/cldm/mmdit.py +++ b/comfy/cldm/mmdit.py @@ -6,6 +6,7 @@ class ControlNet(comfy.ldm.modules.diffusionmodules.mmdit.MMDiT): def __init__( self, num_blocks = None, + control_latent_channels = None, dtype = None, device = None, operations = None, @@ -17,10 +18,13 @@ class ControlNet(comfy.ldm.modules.diffusionmodules.mmdit.MMDiT): for _ in range(len(self.joint_blocks)): self.controlnet_blocks.append(operations.Linear(self.hidden_size, self.hidden_size, device=device, dtype=dtype)) + if control_latent_channels is None: + control_latent_channels = self.in_channels + self.pos_embed_input = comfy.ldm.modules.diffusionmodules.mmdit.PatchEmbed( None, self.patch_size, - self.in_channels, + control_latent_channels, self.hidden_size, bias=True, strict_img_size=False, diff --git a/comfy/controlnet.py b/comfy/controlnet.py index 860891965..1ea00eccf 100644 --- a/comfy/controlnet.py +++ b/comfy/controlnet.py @@ -79,13 +79,19 @@ class ControlBase: self.previous_controlnet = None self.extra_conds = [] self.strength_type = StrengthType.CONSTANT + self.concat_mask = False + self.extra_concat_orig = [] + self.extra_concat = None - def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None): + def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None, extra_concat=[]): self.cond_hint_original = cond_hint self.strength = strength self.timestep_percent_range = timestep_percent_range if self.latent_format is not None: self.vae = vae + self.extra_concat_orig = extra_concat.copy() + if self.concat_mask and len(self.extra_concat_orig) == 0: + self.extra_concat_orig.append(torch.tensor([[[[1.0]]]])) return self def pre_run(self, model, percent_to_timestep_function): @@ -100,9 +106,9 @@ class ControlBase: def cleanup(self): if self.previous_controlnet is not None: self.previous_controlnet.cleanup() - if self.cond_hint is not None: - del self.cond_hint - self.cond_hint = None + + self.cond_hint = None + self.extra_concat = None self.timestep_range = None def get_models(self): @@ -123,6 +129,8 @@ class ControlBase: c.vae = self.vae c.extra_conds = self.extra_conds.copy() c.strength_type = self.strength_type + c.concat_mask = self.concat_mask + c.extra_concat_orig = self.extra_concat_orig.copy() def inference_memory_requirements(self, dtype): if self.previous_controlnet is not None: @@ -175,7 +183,7 @@ class ControlBase: class ControlNet(ControlBase): - def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT): + def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False): super().__init__(device) self.control_model = control_model self.load_device = load_device @@ -189,6 +197,7 @@ class ControlNet(ControlBase): self.latent_format = latent_format self.extra_conds += extra_conds self.strength_type = strength_type + self.concat_mask = concat_mask def get_control(self, x_noisy, t, cond, batched_number): control_prev = None @@ -220,6 +229,13 @@ class ControlNet(ControlBase): comfy.model_management.load_models_gpu(loaded_models) if self.latent_format is not None: self.cond_hint = self.latent_format.process_in(self.cond_hint) + if len(self.extra_concat_orig) > 0: + to_concat = [] + for c in self.extra_concat_orig: + c = comfy.utils.common_upscale(c, self.cond_hint.shape[3], self.cond_hint.shape[2], self.upscale_algorithm, "center") + to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0])) + self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1) + self.cond_hint = self.cond_hint.to(device=self.device, dtype=dtype) if x_noisy.shape[0] != self.cond_hint.shape[0]: self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number) @@ -410,12 +426,17 @@ def load_controlnet_mmdit(sd): for k in sd: new_sd[k] = sd[k] - control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) + concat_mask = False + control_latent_channels = new_sd.get("pos_embed_input.proj.weight").shape[1] + if control_latent_channels == 17: #inpaint controlnet + concat_mask = True + + control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) control_model = controlnet_load_state_dict(control_model, new_sd) latent_format = comfy.latent_formats.SD3() latent_format.shift_factor = 0 #SD3 controlnet weirdness - control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype) + control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype) return control @@ -450,13 +471,16 @@ def load_controlnet_flux_instantx(sd): num_union_modes = new_sd[union_cnet].shape[0] control_latent_channels = new_sd.get("pos_embed_input.weight").shape[1] // 4 + concat_mask = False + if control_latent_channels == 17: + concat_mask = True control_model = comfy.ldm.flux.controlnet.ControlNetFlux(latent_input=True, num_union_modes=num_union_modes, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config) control_model = controlnet_load_state_dict(control_model, new_sd) latent_format = comfy.latent_formats.Flux() extra_conds = ['y', 'guidance'] - 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) + control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds) return control def convert_mistoline(sd): diff --git a/comfy_extras/nodes_controlnet.py b/comfy_extras/nodes_controlnet.py index 0773c8a5b..7cf6ce607 100644 --- a/comfy_extras/nodes_controlnet.py +++ b/comfy_extras/nodes_controlnet.py @@ -1,4 +1,6 @@ from comfy.cldm.control_types import UNION_CONTROLNET_TYPES +import nodes +import comfy.utils class SetUnionControlNetType: @classmethod @@ -22,6 +24,33 @@ class SetUnionControlNetType: return (control_net,) +class ControlNetInpaintingAliMamaApply(nodes.ControlNetApplyAdvanced): + @classmethod + def INPUT_TYPES(s): + return {"required": {"positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "control_net": ("CONTROL_NET", ), + "vae": ("VAE", ), + "image": ("IMAGE", ), + "mask": ("MASK", ), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), + "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}) + }} + + FUNCTION = "apply_inpaint_controlnet" + + CATEGORY = "conditioning/controlnet" + + def apply_inpaint_controlnet(self, positive, negative, control_net, vae, image, mask, strength, start_percent, end_percent): + mask = 1.0 - mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])) + mask_apply = comfy.utils.common_upscale(mask, image.shape[2], image.shape[1], "bilinear", "center").round() + image = image * mask_apply.movedim(1, -1).repeat(1, 1, 1, image.shape[3]) + return self.apply_controlnet(positive, negative, control_net, image, strength, start_percent, end_percent, vae=vae, extra_concat=[mask]) + + + NODE_CLASS_MAPPINGS = { "SetUnionControlNetType": SetUnionControlNetType, + "ControlNetInpaintingAliMamaApply": ControlNetInpaintingAliMamaApply, } diff --git a/nodes.py b/nodes.py index 0d3749d42..84524b6b9 100644 --- a/nodes.py +++ b/nodes.py @@ -824,7 +824,7 @@ class ControlNetApplyAdvanced: CATEGORY = "conditioning/controlnet" - def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent, vae=None): + def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent, vae=None, extra_concat=[]): if strength == 0: return (positive, negative) @@ -841,7 +841,7 @@ class ControlNetApplyAdvanced: if prev_cnet in cnets: c_net = cnets[prev_cnet] else: - c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae) + c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae=vae, extra_concat=extra_concat) c_net.set_previous_controlnet(prev_cnet) cnets[prev_cnet] = c_net From ca08597670c180554ab494e9452e12132d9b346a Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 14 Sep 2024 09:17:13 -0400 Subject: [PATCH 24/26] Make the inpaint controlnet node work with non inpaint ones. --- comfy_extras/nodes_controlnet.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/comfy_extras/nodes_controlnet.py b/comfy_extras/nodes_controlnet.py index 7cf6ce607..2d20e1fed 100644 --- a/comfy_extras/nodes_controlnet.py +++ b/comfy_extras/nodes_controlnet.py @@ -43,10 +43,14 @@ class ControlNetInpaintingAliMamaApply(nodes.ControlNetApplyAdvanced): CATEGORY = "conditioning/controlnet" def apply_inpaint_controlnet(self, positive, negative, control_net, vae, image, mask, strength, start_percent, end_percent): - mask = 1.0 - mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])) - mask_apply = comfy.utils.common_upscale(mask, image.shape[2], image.shape[1], "bilinear", "center").round() - image = image * mask_apply.movedim(1, -1).repeat(1, 1, 1, image.shape[3]) - return self.apply_controlnet(positive, negative, control_net, image, strength, start_percent, end_percent, vae=vae, extra_concat=[mask]) + extra_concat = [] + if control_net.concat_mask: + mask = 1.0 - mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])) + mask_apply = comfy.utils.common_upscale(mask, image.shape[2], image.shape[1], "bilinear", "center").round() + image = image * mask_apply.movedim(1, -1).repeat(1, 1, 1, image.shape[3]) + extra_concat = [mask] + + return self.apply_controlnet(positive, negative, control_net, image, strength, start_percent, end_percent, vae=vae, extra_concat=extra_concat) From 5e68a4ce67450dd8d3aa1c51c6d06014ac46c03f Mon Sep 17 00:00:00 2001 From: JettHu <35261585+JettHu@users.noreply.github.com> Date: Sun, 15 Sep 2024 13:03:09 +0800 Subject: [PATCH 25/26] Reduce repeated calls of INPUT_TYPES in cache (#4922) --- comfy_execution/caching.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index 62311ed7f..630f280fc 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -1,11 +1,21 @@ import itertools -from typing import Sequence, Mapping +from typing import Sequence, Mapping, Dict from comfy_execution.graph import DynamicPrompt import nodes from comfy_execution.graph_utils import is_link +NODE_CLASS_CONTAINS_UNIQUE_ID: Dict[str, bool] = {} + + +def include_unique_id_in_input(class_type: str) -> bool: + if class_type in NODE_CLASS_CONTAINS_UNIQUE_ID: + return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] + class_def = nodes.NODE_CLASS_MAPPINGS[class_type] + NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] = "UNIQUE_ID" in class_def.INPUT_TYPES().get("hidden", {}).values() + return NODE_CLASS_CONTAINS_UNIQUE_ID[class_type] + class CacheKeySet: def __init__(self, dynprompt, node_ids, is_changed_cache): self.keys = {} @@ -98,7 +108,7 @@ class CacheKeySetInputSignature(CacheKeySet): class_type = node["class_type"] class_def = nodes.NODE_CLASS_MAPPINGS[class_type] signature = [class_type, self.is_changed_cache.get(node_id)] - if self.include_node_id_in_input() or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT) or "UNIQUE_ID" in class_def.INPUT_TYPES().get("hidden", {}).values(): + if self.include_node_id_in_input() or (hasattr(class_def, "NOT_IDEMPOTENT") and class_def.NOT_IDEMPOTENT) or include_unique_id_in_input(class_type): signature.append(node_id) inputs = node["inputs"] for key in sorted(inputs.keys()): From e813abbb2c27448735d02af819e79f1a036b7212 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 15 Sep 2024 07:59:18 -0400 Subject: [PATCH 26/26] Long CLIP L support for SDXL, SD3 and Flux. Use the *CLIPLoader nodes. --- comfy/sd.py | 12 +++++------- comfy/sd1_clip.py | 2 ++ comfy/sdxl_clip.py | 9 ++++++--- comfy/text_encoders/flux.py | 6 ++++-- comfy/text_encoders/long_clipl.py | 15 +++++++++++++-- comfy/text_encoders/sd3_clip.py | 9 ++++++--- 6 files changed, 36 insertions(+), 17 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index 1e1a6594b..07310b9d4 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -445,12 +445,8 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip clip_target.tokenizer = comfy.text_encoders.sa_t5.SAT5Tokenizer else: w = clip_data[0].get("text_model.embeddings.position_embedding.weight", None) - if w is not None and w.shape[0] == 248: - clip_target.clip = comfy.text_encoders.long_clipl.LongClipModel - clip_target.tokenizer = comfy.text_encoders.long_clipl.LongClipTokenizer - else: - clip_target.clip = sd1_clip.SD1ClipModel - clip_target.tokenizer = sd1_clip.SD1Tokenizer + clip_target.clip = sd1_clip.SD1ClipModel + clip_target.tokenizer = sd1_clip.SD1Tokenizer elif len(clip_data) == 2: if clip_type == CLIPType.SD3: clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=True, clip_g=True, t5=False) @@ -475,10 +471,12 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer parameters = 0 + tokenizer_data = {} for c in clip_data: parameters += comfy.utils.calculate_parameters(c) + tokenizer_data, model_options = comfy.text_encoders.long_clipl.model_options_long_clip(c, tokenizer_data, model_options) - clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters, model_options=model_options) + clip = CLIP(clip_target, embedding_directory=embedding_directory, parameters=parameters, tokenizer_data=tokenizer_data, model_options=model_options) for c in clip_data: m, u = clip.load_sd(c) if len(m) > 0: diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py index 676653f77..9f0d95b0a 100644 --- a/comfy/sd1_clip.py +++ b/comfy/sd1_clip.py @@ -542,6 +542,7 @@ class SD1Tokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}, clip_name="l", tokenizer=SDTokenizer): self.clip_name = clip_name self.clip = "clip_{}".format(self.clip_name) + tokenizer = tokenizer_data.get("{}_tokenizer_class".format(self.clip), tokenizer) setattr(self, self.clip, tokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)) def tokenize_with_weights(self, text:str, return_word_ids=False): @@ -570,6 +571,7 @@ class SD1ClipModel(torch.nn.Module): self.clip_name = clip_name self.clip = "clip_{}".format(self.clip_name) + clip_model = model_options.get("{}_class".format(self.clip), clip_model) setattr(self, self.clip, clip_model(device=device, dtype=dtype, model_options=model_options, **kwargs)) self.dtypes = set() diff --git a/comfy/sdxl_clip.py b/comfy/sdxl_clip.py index a0145caa4..4d0a4e8e7 100644 --- a/comfy/sdxl_clip.py +++ b/comfy/sdxl_clip.py @@ -22,7 +22,8 @@ class SDXLClipGTokenizer(sd1_clip.SDTokenizer): class SDXLTokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}): - self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory) + clip_l_tokenizer_class = tokenizer_data.get("clip_l_tokenizer_class", sd1_clip.SDTokenizer) + self.clip_l = clip_l_tokenizer_class(embedding_directory=embedding_directory) self.clip_g = SDXLClipGTokenizer(embedding_directory=embedding_directory) def tokenize_with_weights(self, text:str, return_word_ids=False): @@ -40,7 +41,8 @@ class SDXLTokenizer: class SDXLClipModel(torch.nn.Module): def __init__(self, device="cpu", dtype=None, model_options={}): super().__init__() - self.clip_l = sd1_clip.SDClipModel(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, model_options=model_options) + clip_l_class = model_options.get("clip_l_class", sd1_clip.SDClipModel) + self.clip_l = clip_l_class(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, model_options=model_options) self.clip_g = SDXLClipG(device=device, dtype=dtype, model_options=model_options) self.dtypes = set([dtype]) @@ -57,7 +59,8 @@ class SDXLClipModel(torch.nn.Module): token_weight_pairs_l = token_weight_pairs["l"] g_out, g_pooled = self.clip_g.encode_token_weights(token_weight_pairs_g) l_out, l_pooled = self.clip_l.encode_token_weights(token_weight_pairs_l) - return torch.cat([l_out, g_out], dim=-1), g_pooled + cut_to = min(l_out.shape[1], g_out.shape[1]) + return torch.cat([l_out[:,:cut_to], g_out[:,:cut_to]], dim=-1), g_pooled def load_sd(self, sd): if "text_model.encoder.layers.30.mlp.fc1.weight" in sd: diff --git a/comfy/text_encoders/flux.py b/comfy/text_encoders/flux.py index 91d1f249d..7c75fe64b 100644 --- a/comfy/text_encoders/flux.py +++ b/comfy/text_encoders/flux.py @@ -18,7 +18,8 @@ class T5XXLTokenizer(sd1_clip.SDTokenizer): class FluxTokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}): - self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory) + clip_l_tokenizer_class = tokenizer_data.get("clip_l_tokenizer_class", sd1_clip.SDTokenizer) + self.clip_l = clip_l_tokenizer_class(embedding_directory=embedding_directory) self.t5xxl = T5XXLTokenizer(embedding_directory=embedding_directory) def tokenize_with_weights(self, text:str, return_word_ids=False): @@ -38,7 +39,8 @@ class FluxClipModel(torch.nn.Module): def __init__(self, dtype_t5=None, device="cpu", dtype=None, model_options={}): super().__init__() dtype_t5 = comfy.model_management.pick_weight_dtype(dtype_t5, dtype, device) - self.clip_l = sd1_clip.SDClipModel(device=device, dtype=dtype, return_projected_pooled=False, model_options=model_options) + clip_l_class = model_options.get("clip_l_class", sd1_clip.SDClipModel) + self.clip_l = clip_l_class(device=device, dtype=dtype, return_projected_pooled=False, model_options=model_options) self.t5xxl = T5XXLModel(device=device, dtype=dtype_t5, model_options=model_options) self.dtypes = set([dtype, dtype_t5]) diff --git a/comfy/text_encoders/long_clipl.py b/comfy/text_encoders/long_clipl.py index 4677fb3b0..b81912cb3 100644 --- a/comfy/text_encoders/long_clipl.py +++ b/comfy/text_encoders/long_clipl.py @@ -6,9 +6,9 @@ class LongClipTokenizer_(sd1_clip.SDTokenizer): super().__init__(max_length=248, embedding_directory=embedding_directory, tokenizer_data=tokenizer_data) class LongClipModel_(sd1_clip.SDClipModel): - def __init__(self, device="cpu", dtype=None, model_options={}): + def __init__(self, *args, **kwargs): textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "long_clipl.json") - super().__init__(device=device, textmodel_json_config=textmodel_json_config, return_projected_pooled=False, dtype=dtype, model_options=model_options) + super().__init__(*args, textmodel_json_config=textmodel_json_config, **kwargs) class LongClipTokenizer(sd1_clip.SD1Tokenizer): def __init__(self, embedding_directory=None, tokenizer_data={}): @@ -17,3 +17,14 @@ class LongClipTokenizer(sd1_clip.SD1Tokenizer): class LongClipModel(sd1_clip.SD1ClipModel): def __init__(self, device="cpu", dtype=None, model_options={}, **kwargs): super().__init__(device=device, dtype=dtype, model_options=model_options, clip_model=LongClipModel_, **kwargs) + +def model_options_long_clip(sd, tokenizer_data, model_options): + w = sd.get("clip_l.text_model.embeddings.position_embedding.weight", None) + if w is None: + w = sd.get("text_model.embeddings.position_embedding.weight", None) + if w is not None and w.shape[0] == 248: + tokenizer_data = tokenizer_data.copy() + model_options = model_options.copy() + tokenizer_data["clip_l_tokenizer_class"] = LongClipTokenizer_ + model_options["clip_l_class"] = LongClipModel_ + return tokenizer_data, model_options diff --git a/comfy/text_encoders/sd3_clip.py b/comfy/text_encoders/sd3_clip.py index e3832ac24..c54f28856 100644 --- a/comfy/text_encoders/sd3_clip.py +++ b/comfy/text_encoders/sd3_clip.py @@ -20,7 +20,8 @@ class T5XXLTokenizer(sd1_clip.SDTokenizer): class SD3Tokenizer: def __init__(self, embedding_directory=None, tokenizer_data={}): - self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory) + clip_l_tokenizer_class = tokenizer_data.get("clip_l_tokenizer_class", sd1_clip.SDTokenizer) + self.clip_l = clip_l_tokenizer_class(embedding_directory=embedding_directory) self.clip_g = sdxl_clip.SDXLClipGTokenizer(embedding_directory=embedding_directory) self.t5xxl = T5XXLTokenizer(embedding_directory=embedding_directory) @@ -42,7 +43,8 @@ class SD3ClipModel(torch.nn.Module): super().__init__() self.dtypes = set() if clip_l: - self.clip_l = sd1_clip.SDClipModel(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, return_projected_pooled=False, model_options=model_options) + clip_l_class = model_options.get("clip_l_class", sd1_clip.SDClipModel) + self.clip_l = clip_l_class(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False, return_projected_pooled=False, model_options=model_options) self.dtypes.add(dtype) else: self.clip_l = None @@ -95,7 +97,8 @@ class SD3ClipModel(torch.nn.Module): if self.clip_g is not None: g_out, g_pooled = self.clip_g.encode_token_weights(token_weight_pairs_g) if lg_out is not None: - lg_out = torch.cat([lg_out, g_out], dim=-1) + cut_to = min(lg_out.shape[1], g_out.shape[1]) + lg_out = torch.cat([lg_out[:,:cut_to], g_out[:,:cut_to]], dim=-1) else: lg_out = torch.nn.functional.pad(g_out, (768, 0)) else: