Merge remote-tracking branch 'origin/main' into feat/model-downloader-tokens
# Conflicts: # CHANGELOG.md
This commit is contained in:
Vito Sansevero
@@ -2,6 +2,14 @@
|
||||
|
||||
## 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.
|
||||
|
||||
### Model Downloader
|
||||
|
||||
- Civitai tokens and Hugging Face tokens are read on the ComfyUI server:
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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];
|
||||
|
||||
Reference in New Issue
Block a user