diff --git a/app/user_manager.py b/app/user_manager.py index 62c22cde5..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: @@ -120,23 +120,44 @@ 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" - 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).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).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: - results = [[x] + x.split(os.sep) for x in results] + if split_path and not full_info: + results = [[x] + x.split('/') for x in results] return web.json_response(results) @@ -144,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}") @@ -159,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}") @@ -167,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) @@ -176,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) @@ -187,7 +208,7 @@ class UserManager(): return path os.remove(path) - + return web.Response(status=204) @routes.post("/userdata/{file}/move/{dest}") @@ -195,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) 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) 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 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..936c6bd27 --- /dev/null +++ b/tests-unit/prompt_server_test/user_manager_test.py @@ -0,0 +1,120 @@ +import pytest +import os +from aiohttp import web +from app.user_manager import UserManager +from unittest.mock import patch + +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", "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() == [ + ["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 + + +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"