From 9be3cc3397e3b29b817206dd854da6e64c391304 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 7 Nov 2026 03:54:23 -0800 Subject: [PATCH 1/2] feat(local_image_loader): add a folder output for the browsed folder The Local Image Loader can browse folders visually, but only output a picked image. Add a third output, folder: the path of the folder the gallery shows, or the picked image's folder until the user browses. Browsing is enough, so a training dataset can be found by its thumbnails and passed to a node that reads a whole folder. The gallery stores the shown folder through set_node_selection (type "folder"), which accepts only directories inside the allowed roots; load_media checks the stored folder against the roots again. --- CHANGELOG.md | 10 +++ README.md | 2 + examples/documentation/local_image_loader.md | 5 +- kikotools/tools/local_image_loader/node.py | 36 +++++++-- .../tools/local_image_loader/test_routes.py | 34 +++++++-- tests/unit/tools/test_local_image_loader.py | 75 +++++++++++++++++-- web/kiko_local_image_loader.js | 23 ++++++ 7 files changed, 165 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 94d6745..65fcc30 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Changelog +## Unreleased + +### Local Image Loader + +- New `folder` output: the path of the folder the gallery shows, or the + selected image's folder until you browse. Browsing to a folder is enough; + no image needs to be selected. Use it to pick a training dataset visually and + pass it to a node that reads a whole folder. The folder must be inside a + saved path or ComfyUI's input/output, as for images. + ## 1.3.1 ### Startup diff --git a/README.md b/README.md index 9faae93..174de41 100644 --- a/README.md +++ b/README.md @@ -427,6 +427,7 @@ Visual gallery browser for loading local images, videos, and audio files directl - **Double-Click Preview**: Open full-size media in new browser tab - **Smart Sorting**: Sort by name, date, or file size in ascending or descending order - **Pagination Support**: Efficiently browse large directories with page controls +- **Folder Output**: Outputs the path of the folder the gallery shows, so a folder can be picked by browsing **Use Cases:** - Load reference images from local folders for img2img workflows @@ -434,6 +435,7 @@ Visual gallery browser for loading local images, videos, and audio files directl - Quickly access frequently used asset directories - Extract prompts and settings from previously generated images - Preview media files before loading into workflow +- Browse to a training dataset and pass its folder to a captioning or training node #### 🌐 Model Downloader Download models, LoRAs, and other assets directly from CivitAI, HuggingFace, and custom URLs within ComfyUI. diff --git a/examples/documentation/local_image_loader.md b/examples/documentation/local_image_loader.md index befb56e..755efc8 100644 --- a/examples/documentation/local_image_loader.md +++ b/examples/documentation/local_image_loader.md @@ -27,9 +27,10 @@ None - The node uses a visual interface for file selection | Output | Type | Description | |--------|------|-------------| | `image` | IMAGE | The selected image as a tensor | -| `video_path` | STRING | Path to the selected video file | -| `audio_path` | STRING | Path to the selected audio file | | `info` | STRING | JSON metadata about the selected image | +| `folder` | STRING | The folder the gallery is showing, or the selected image's folder | + +Browse to a folder and queue: `folder` passes its path on without selecting an image, for nodes that read a whole folder, such as a dataset captioner. ## Usage diff --git a/kikotools/tools/local_image_loader/node.py b/kikotools/tools/local_image_loader/node.py index eacdb32..114521f 100644 --- a/kikotools/tools/local_image_loader/node.py +++ b/kikotools/tools/local_image_loader/node.py @@ -113,10 +113,17 @@ class LocalImageLoaderNode(ComfyAssetsBaseNode): RETURN_TYPES = ( "IMAGE", "STRING", + "STRING", ) RETURN_NAMES = ( "image", "info", + "folder", + ) + OUTPUT_TOOLTIPS = ( + "The selected image.", + "The selected image's metadata as JSON.", + "The folder the gallery shows, e.g. a training dataset to pass on by path.", ) FUNCTION = "load_media" CATEGORY = "🫶 ComfyAssets/💾 Images" @@ -128,7 +135,7 @@ class LocalImageLoaderNode(ComfyAssetsBaseNode): return os.path.getmtime(SELECTIONS_FILE) return float("inf") - def load_media(self, unique_id: str) -> Tuple[torch.Tensor, str]: + def load_media(self, unique_id: str) -> Tuple[torch.Tensor, str, str]: """ Load selected media based on node's unique ID. @@ -136,13 +143,15 @@ class LocalImageLoaderNode(ComfyAssetsBaseNode): unique_id: Unique identifier for this node instance Returns: - Tuple of (image tensor, info string) + Tuple of (image tensor, info string, folder path) """ image_tensor = create_empty_tensor() info_string = "" + folder = "" selections = load_selections() node_selections = selections.get(str(unique_id), {}) + roots = get_allowed_roots() # Load image if selected image_selection = node_selections.get("image") @@ -150,17 +159,27 @@ class LocalImageLoaderNode(ComfyAssetsBaseNode): # Load the resolved path we validated, so a symlink swapped after # the check can't redirect the read. image_path = resolve_path(image_selection["path"]) - if not is_allowed_image_path(image_path, get_allowed_roots()): + if not is_allowed_image_path(image_path, roots): # selections.json may have been tampered with logger.warning("Refusing to load disallowed path: %s", image_path) elif os.path.exists(image_path): try: image_tensor, metadata = load_image_from_path(image_path) info_string = json.dumps(metadata, indent=4, ensure_ascii=False) + folder = os.path.dirname(image_path) except Exception as e: logger.error("Error loading image: %s", e) - return (image_tensor, info_string) + # The folder the gallery shows wins over the picked image's folder, so + # browsing to a dataset is enough to pass it on. + folder_selection = node_selections.get("folder") + browsed = resolve_path( + folder_selection.get("path") if isinstance(folder_selection, dict) else None + ) + if browsed and os.path.isdir(browsed) and is_within_roots(browsed, roots): + folder = browsed + + return (image_tensor, info_string, folder) # --------------------------------------------------------------------------- @@ -215,8 +234,13 @@ async def set_node_selection(request): ): return _error(400, "Missing required data.") - # The frontend also selects video/audio; load_media only ever reads images. - if not is_allowed_media_path(path, get_allowed_roots(), MEDIA_KINDS): + if media_type == "folder": + # The folder the gallery is showing, output as a path. + allowed = os.path.isdir(path) and is_within_roots(path, get_allowed_roots()) + else: + # The frontend also selects video/audio; load_media only ever reads images. + allowed = is_allowed_media_path(path, get_allowed_roots(), MEDIA_KINDS) + if not allowed: return _error(403, "Path not allowed.") selections = load_selections() diff --git a/tests/unit/tools/local_image_loader/test_routes.py b/tests/unit/tools/local_image_loader/test_routes.py index baeaf9a..6b01872 100644 --- a/tests/unit/tools/local_image_loader/test_routes.py +++ b/tests/unit/tools/local_image_loader/test_routes.py @@ -156,6 +156,25 @@ class TestViewMedia: r = run(node_mod.get_thumbnail(req({"filepath": str(env.root / "clip.mp4")}))) assert r.status == 403 + def test_folder_selection_ok(self, env): + p = str(env.root / "sub") + r = run( + node_mod.set_node_selection( + req(body={"node_id": 2, "path": p, "type": "folder"}) + ) + ) + assert r.status == 200 + assert json.loads(env.sel.read_text())["2"]["folder"] == {"path": p} + + def test_folder_selection_disallowed_403(self, env): + for p in (str(env.outside), str(env.root / "a.png"), "/etc"): + r = run( + node_mod.set_node_selection( + req(body={"node_id": 2, "path": p, "type": "folder"}) + ) + ) + assert r.status == 403, p + def test_video_selection_allowed_but_not_loaded(self, env): (env.root / "clip.mp4").write_bytes(b"x") p = str(env.root / "clip.mp4") @@ -380,31 +399,32 @@ class TestLoadMediaGuard: env.sel.write_text( json.dumps({"7": {"image": {"path": str(env.outside / "o.png")}}}) ) - image, info = node_mod.LocalImageLoaderNode().load_media("7") - assert image.shape == (1, 1, 1, 4) and info == "" + image, info, folder = node_mod.LocalImageLoaderNode().load_media("7") + assert image.shape == (1, 1, 1, 4) and info == "" and folder == "" def test_non_image_selection_refused(self, env): env.sel.write_text( json.dumps({"7": {"image": {"path": str(env.root / "notes.txt")}}}) ) - _, info = node_mod.LocalImageLoaderNode().load_media("7") - assert info == "" + _, info, folder = node_mod.LocalImageLoaderNode().load_media("7") + assert info == "" and folder == "" def test_allowed_selection_loads(self, env): env.sel.write_text( json.dumps({"7": {"image": {"path": str(env.root / "a.png")}}}) ) - image, info = node_mod.LocalImageLoaderNode().load_media("7") + image, info, folder = node_mod.LocalImageLoaderNode().load_media("7") assert image.shape == (1, 300, 400, 3) assert json.loads(info)["width"] == 400 + assert folder == os.path.realpath(env.root) def test_load_error_is_swallowed(self, env): (env.root / "bad.png").write_bytes(b"junk") env.sel.write_text( json.dumps({"7": {"image": {"path": str(env.root / "bad.png")}}}) ) - image, info = node_mod.LocalImageLoaderNode().load_media("7") - assert info == "" + image, info, folder = node_mod.LocalImageLoaderNode().load_media("7") + assert info == "" and folder == "" class TestPersistence: diff --git a/tests/unit/tools/test_local_image_loader.py b/tests/unit/tools/test_local_image_loader.py index 9e509e1..c30e2b8 100644 --- a/tests/unit/tools/test_local_image_loader.py +++ b/tests/unit/tools/test_local_image_loader.py @@ -201,8 +201,8 @@ class TestLocalImageLoaderNode: def test_node_properties(self): """Test node properties.""" - assert LocalImageLoaderNode.RETURN_TYPES == ("IMAGE", "STRING") - assert LocalImageLoaderNode.RETURN_NAMES == ("image", "info") + assert LocalImageLoaderNode.RETURN_TYPES == ("IMAGE", "STRING", "STRING") + assert LocalImageLoaderNode.RETURN_NAMES == ("image", "info", "folder") assert LocalImageLoaderNode.FUNCTION == "load_media" assert LocalImageLoaderNode.CATEGORY == "🫶 ComfyAssets/💾 Images" @@ -212,13 +212,14 @@ class TestLocalImageLoaderNode: mock_load_selections.return_value = {} node = LocalImageLoaderNode() - image, info = node.load_media("test_id") + image, info, folder = node.load_media("test_id") # Check empty returns assert isinstance(image, torch.Tensor) assert image.shape == (1, 1, 1, 4) assert torch.all(image == 0) assert info == "" + assert folder == "" @patch("kikotools.tools.local_image_loader.node.load_selections") @patch("kikotools.tools.local_image_loader.node.load_image_from_path") @@ -244,11 +245,12 @@ class TestLocalImageLoaderNode: ), ): node = LocalImageLoaderNode() - image, info = node.load_media("test_id") + image, info, folder = node.load_media("test_id") # Check returns assert torch.equal(image, test_tensor) assert json.loads(info) == test_metadata + assert folder == os.path.realpath("/path/to") @patch("kikotools.tools.local_image_loader.node.load_selections") def test_load_media_ignores_video_audio_selection(self, mock_load_selections): @@ -260,10 +262,73 @@ class TestLocalImageLoaderNode: } } - image, info = LocalImageLoaderNode().load_media("test_id") + image, info, folder = LocalImageLoaderNode().load_media("test_id") assert image.shape == (1, 1, 1, 4) assert info == "" + assert folder == "" + + def test_load_media_outputs_browsed_folder(self, tmp_path): + """The folder the gallery is showing is output even with no image picked.""" + (tmp_path / "dataset").mkdir() + folder_path = str(tmp_path / "dataset") + with ( + patch( + "kikotools.tools.local_image_loader.node.load_selections", + return_value={"test_id": {"folder": {"path": folder_path}}}, + ), + patch( + "kikotools.tools.local_image_loader.node.get_allowed_roots", + return_value=[str(tmp_path)], + ), + ): + image, info, folder = LocalImageLoaderNode().load_media("test_id") + + assert image.shape == (1, 1, 1, 4) + assert folder == os.path.realpath(folder_path) + + def test_load_media_browsed_folder_wins_over_image_folder(self, tmp_path): + """Browsing away after picking an image outputs the folder now shown.""" + for d in ("a", "b"): + (tmp_path / d).mkdir() + Image.new("RGB", (4, 4)).save(tmp_path / "a" / "x.png") + selection = { + "image": {"path": str(tmp_path / "a" / "x.png")}, + "folder": {"path": str(tmp_path / "b")}, + } + with ( + patch( + "kikotools.tools.local_image_loader.node.load_selections", + return_value={"test_id": selection}, + ), + patch( + "kikotools.tools.local_image_loader.node.get_allowed_roots", + return_value=[str(tmp_path)], + ), + ): + _, _, folder = LocalImageLoaderNode().load_media("test_id") + + assert folder == os.path.realpath(tmp_path / "b") + + def test_load_media_refuses_folder_outside_roots(self, tmp_path): + """A tampered selections.json cannot point the folder output elsewhere.""" + (tmp_path / "root").mkdir() + (tmp_path / "elsewhere").mkdir() + with ( + patch( + "kikotools.tools.local_image_loader.node.load_selections", + return_value={ + "test_id": {"folder": {"path": str(tmp_path / "elsewhere")}} + }, + ), + patch( + "kikotools.tools.local_image_loader.node.get_allowed_roots", + return_value=[str(tmp_path / "root")], + ), + ): + _, _, folder = LocalImageLoaderNode().load_media("test_id") + + assert folder == "" def test_is_changed(self): """Test IS_CHANGED method.""" diff --git a/web/kiko_local_image_loader.js b/web/kiko_local_image_loader.js index bf8f23b..8f8d34a 100644 --- a/web/kiko_local_image_loader.js +++ b/web/kiko_local_image_loader.js @@ -436,6 +436,7 @@ app.registerExtension({ pathInput.value = currentDirectory; upFolderBtn.disabled = !parentDirectory; + selectFolder(currentDirectory); renderGallery(); renderPagination(); @@ -575,6 +576,28 @@ app.registerExtension({ } // Select media + // The folder output is the folder the gallery shows. + let selectedFolder = null; + async function selectFolder(directory) { + if (directory === selectedFolder) return; + selectedFolder = directory; + try { + const response = await fetchApi('/kiko_local_image_loader/set_node_selection', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ node_id: node.id, path: directory, type: 'folder' }) + }); + if (!response.ok) { + selectedFolder = null; + showStatus(`Folder not set: ${await describeHttpError(response)}`); + } + } catch (error) { + if (scope.signal.aborted) return; + selectedFolder = null; + console.error('Error setting folder:', error); + } + } + async function selectMedia(item, type, cardElement) { // Update selection const previousPath = selectedPaths[type]; From 508eefddfd519287f6af5ce5a26509b82b6a5e36 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 7 Nov 2026 04:06:25 -0800 Subject: [PATCH 2/2] test(local_image_loader): cover the folder output in the browser suite The test host unpacked two values from load_media, and the gallery test expected no stored selection although the gallery now stores the folder it shows. Return and assert the folder instead. --- tests/browser/server.py | 8 ++++++-- tests/browser/workflows.spec.mjs | 9 ++++++--- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/browser/server.py b/tests/browser/server.py index 397c03e..90c30d4 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -118,9 +118,13 @@ def create_app(root: Path) -> web.Application: cls = plugin.NODE_CLASS_MAPPINGS[name] result = await asyncio.to_thread(getattr(cls(), cls.FUNCTION), **data) if name == "KikoLocalImageLoader": - tensor, info = result + tensor, info, folder = result return web.json_response( - {"shape": list(tensor.shape), "info": json.loads(info) if info else {}} + { + "shape": list(tensor.shape), + "info": json.loads(info) if info else {}, + "folder": folder, + } ) return web.json_response(result) diff --git a/tests/browser/workflows.spec.mjs b/tests/browser/workflows.spec.mjs index 5f08f4d..b4338e4 100644 --- a/tests/browser/workflows.spec.mjs +++ b/tests/browser/workflows.spec.mjs @@ -225,14 +225,16 @@ test('local image selection persists and the real node loads the selected PNG te const image = page.getByAltText('cat & dog.png', { exact: true }); await expect(image).toBeVisible(); await expect.poll(() => image.evaluate(image => image.complete && image.naturalWidth > 0)).toBe(true); - const selected = page.waitForResponse(response => response.url().endsWith('/set_node_selection')); + const selected = page.waitForResponse(response => response.url().endsWith('/set_node_selection') && response.request().postDataJSON().type === 'image'); await image.click(); expect((await selected).ok()).toBe(true); const selectedPath = `${(await state(request)).input}/cat & dog.png`; expect((await state(request)).selections[String(id)].image.path).toBe(selectedPath); const executed = await request.post('/__test__/execute/KikoLocalImageLoader', { data: { unique_id: String(id) } }); await expect(executed).toBeOK(); - expect((await executed.json()).shape).toEqual([1, 12, 16, 3]); + const output = await executed.json(); + expect(output.shape).toEqual([1, 12, 16, 3]); + expect(output.folder).toBe((await state(request)).input); await page.reload(); await expect(page.locator('body')).toHaveAttribute('data-ready', 'true'); expect(await createNode(page, 'KikoLocalImageLoader')).toBe(id); @@ -271,7 +273,8 @@ test('gallery rejects disallowed folders and retries after a network error', asy await expect(page.locator('.kiko-status-message')).toBeVisible(); const deniedSelection = await request.post('/kiko_local_image_loader/set_node_selection', { data: { node_id: id, path: `${info.outside}/secret.png`, type: 'image' } }); expect(deniedSelection.status()).toBe(403); - expect((await state(request)).selections).toEqual({}); + // Only the folder the gallery still shows is stored, never the denied one. + expect((await state(request)).selections).toEqual({ [String(id)]: { folder: { path: info.input } } }); await page.locator('.kiko-path-input').fill(info.input); await page.route('**/kiko_local_image_loader/images?*', route => route.abort('failed'), { times: 1 }); await page.locator('.kiko-browse').click();