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 1/7] 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 2/7] 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 3/7] 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 4/7] 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 5/7] 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 6/7] 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 7/7] 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)