diff --git a/kikotools/tools/local_image_loader/__init__.py b/kikotools/tools/local_image_loader/__init__.py new file mode 100644 index 0000000..8fde33a --- /dev/null +++ b/kikotools/tools/local_image_loader/__init__.py @@ -0,0 +1,13 @@ +"""Local Image Loader tool for KikoTools.""" + +from .node import LocalImageLoaderNode + +NODE_CLASS_MAPPINGS = { + "KikoLocalImageLoader": LocalImageLoaderNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "KikoLocalImageLoader": "Local Image Loader", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/kikotools/tools/local_image_loader/config.json b/kikotools/tools/local_image_loader/config.json new file mode 100644 index 0000000..6547c30 --- /dev/null +++ b/kikotools/tools/local_image_loader/config.json @@ -0,0 +1,7 @@ +{ + "last_path": "/home/vito/ai-apps/ComfyUI-3.12/output", + "saved_paths": [ + "/home/vito/ai-apps/ComfyUI-3.12/output/2025-05-01", + "/home/vito/ai-apps/ComfyUI-3.12/output/" + ] +} \ No newline at end of file diff --git a/kikotools/tools/local_image_loader/logic.py b/kikotools/tools/local_image_loader/logic.py new file mode 100644 index 0000000..9d1e6fd --- /dev/null +++ b/kikotools/tools/local_image_loader/logic.py @@ -0,0 +1,144 @@ +"""Core logic for Local Image Loader.""" + +import os +import json +import torch +import numpy as np +from PIL import Image +from typing import Tuple, Dict, Any, List + + +def get_supported_extensions() -> Dict[str, List[str]]: + """Get supported file extensions by type.""" + return { + "image": [".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp"], + "video": [".mp4", ".webm", ".mov", ".mkv", ".avi"], + "audio": [".mp3", ".wav", ".ogg", ".flac"], + } + + +def load_image_from_path(path: str) -> Tuple[torch.Tensor, Dict[str, Any]]: + """ + Load an image from the given path and convert it to a tensor. + + Args: + path: Path to the image file + + Returns: + Tuple of (image tensor, metadata dict) + """ + if not os.path.exists(path): + raise FileNotFoundError(f"File not found: {path}") + + with Image.open(path) as img: + # Convert to appropriate format + if "A" in img.getbands(): + img_out = img.convert("RGBA") + else: + img_out = img.convert("RGB") + + # Convert to tensor + img_array = np.array(img_out).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(img_array)[None,] + + # Collect metadata + metadata = { + "filename": os.path.basename(path), + "width": img.width, + "height": img.height, + "mode": img.mode, + "format": img.format, + } + + # Check for embedded metadata + if "parameters" in img.info: + metadata["parameters"] = img.info["parameters"] + if "prompt" in img.info: + try: + metadata["prompt"] = json.loads(img.info["prompt"]) + except (json.JSONDecodeError, TypeError): + metadata["prompt"] = img.info["prompt"] + if "workflow" in img.info: + try: + metadata["workflow"] = json.loads(img.info["workflow"]) + except (json.JSONDecodeError, TypeError): + metadata["workflow"] = img.info["workflow"] + + return image_tensor, metadata + + +def scan_directory( + directory: str, + show_videos: bool = False, + show_audio: bool = False, + sort_by: str = "name", + sort_order: str = "asc", +) -> List[Dict[str, Any]]: + """ + Scan a directory for supported media files. + + Args: + directory: Directory path to scan + show_videos: Include video files + show_audio: Include audio files + sort_by: Sort criteria ('name', 'date', 'size') + sort_order: Sort order ('asc', 'desc') + + Returns: + List of file information dictionaries + """ + if not os.path.isdir(directory): + raise NotADirectoryError(f"Not a directory: {directory}") + + extensions = get_supported_extensions() + items = [] + + for item in os.listdir(directory): + full_path = os.path.join(directory, item) + + try: + stats = os.stat(full_path) + item_data = { + "path": full_path, + "name": item, + "mtime": stats.st_mtime, + "size": stats.st_size, + } + + if os.path.isdir(full_path): + items.append({**item_data, "type": "dir"}) + else: + ext = os.path.splitext(item)[1].lower() + item_type = None + + if ext in extensions["image"]: + item_type = "image" + elif show_videos and ext in extensions["video"]: + item_type = "video" + elif show_audio and ext in extensions["audio"]: + item_type = "audio" + + if item_type: + items.append({**item_data, "type": item_type}) + + except (PermissionError, FileNotFoundError): + continue + + # Sort items + reverse = sort_order == "desc" + if sort_by == "date": + items.sort(key=lambda x: x["mtime"], reverse=reverse) + elif sort_by == "size": + items.sort(key=lambda x: x.get("size", 0), reverse=reverse) + else: # name + items.sort(key=lambda x: x["name"].lower(), reverse=reverse) + + # Directories first + items.sort(key=lambda x: x["type"] != "dir") + + return items + + +def create_empty_tensor() -> torch.Tensor: + """Create an empty tensor for when no image is selected.""" + return torch.zeros(1, 1, 1, 4) diff --git a/kikotools/tools/local_image_loader/node.py b/kikotools/tools/local_image_loader/node.py new file mode 100644 index 0000000..814c2c7 --- /dev/null +++ b/kikotools/tools/local_image_loader/node.py @@ -0,0 +1,291 @@ +"""Local Image Loader node for ComfyUI.""" + +import os +import json +import torch +from typing import Dict, Any, Tuple + +from ...base.base_node import ComfyAssetsBaseNode +from .logic import load_image_from_path, create_empty_tensor + + +NODE_DIR = os.path.dirname(os.path.abspath(__file__)) +SELECTIONS_FILE = os.path.join(NODE_DIR, "selections.json") +CONFIG_FILE = os.path.join(NODE_DIR, "config.json") + + +def load_selections() -> Dict[str, Any]: + """Load node selections from file.""" + if not os.path.exists(SELECTIONS_FILE): + return {} + try: + with open(SELECTIONS_FILE, "r", encoding="utf-8") as f: + return json.load(f) + except (json.JSONDecodeError, IOError): + return {} + + +def save_selections(data: Dict[str, Any]) -> None: + """Save node selections to file.""" + try: + with open(SELECTIONS_FILE, "w", encoding="utf-8") as f: + json.dump(data, f, indent=4, ensure_ascii=False) + except Exception as e: + print(f"KikoLocalImageLoader: Error saving selections: {e}") + + +def load_config() -> Dict[str, Any]: + """Load configuration from file.""" + if os.path.exists(CONFIG_FILE): + try: + with open(CONFIG_FILE, "r", encoding="utf-8") as f: + return json.load(f) + except (json.JSONDecodeError, IOError): + pass + return {} + + +def save_config(data: Dict[str, Any]) -> None: + """Save configuration to file.""" + try: + with open(CONFIG_FILE, "w", encoding="utf-8") as f: + json.dump(data, f, indent=4) + except Exception as e: + print(f"KikoLocalImageLoader: Error saving config: {e}") + + +class LocalImageLoaderNode(ComfyAssetsBaseNode): + """Node for loading images from local filesystem with a visual gallery interface.""" + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """Define input types for the node.""" + return { + "required": {}, + "hidden": {"unique_id": "UNIQUE_ID"}, + } + + RETURN_TYPES = ( + "IMAGE", + "STRING", + "STRING", + "STRING", + ) + RETURN_NAMES = ( + "image", + "video_path", + "audio_path", + "info", + ) + FUNCTION = "load_media" + CATEGORY = "🫶 ComfyAssets/💾 Images" + + @classmethod + def IS_CHANGED(cls, **kwargs): + """Check if node state has changed.""" + if os.path.exists(SELECTIONS_FILE): + return os.path.getmtime(SELECTIONS_FILE) + return float("inf") + + def load_media(self, unique_id: str) -> Tuple[torch.Tensor, str, str, str]: + """ + Load selected media based on node's unique ID. + + Args: + unique_id: Unique identifier for this node instance + + Returns: + Tuple of (image tensor, video path, audio path, info string) + """ + image_tensor = create_empty_tensor() + video_path = "" + audio_path = "" + info_string = "" + + selections = load_selections() + node_selections = selections.get(str(unique_id), {}) + + # Load image if selected + image_selection = node_selections.get("image") + if image_selection and image_selection.get("path"): + image_path = image_selection["path"] + if 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) + except Exception as e: + print(f"KikoLocalImageLoader: Error loading image: {e}") + + # Get video path if selected + video_selection = node_selections.get("video") + if video_selection and video_selection.get("path"): + if os.path.exists(video_selection["path"]): + video_path = video_selection["path"] + + # Get audio path if selected + audio_selection = node_selections.get("audio") + if audio_selection and audio_selection.get("path"): + if os.path.exists(audio_selection["path"]): + audio_path = audio_selection["path"] + + return (image_tensor, video_path, audio_path, info_string) + + +# Setup API routes +try: + import server + from aiohttp import web + import urllib.parse + import io + from PIL import Image + from .logic import scan_directory + + prompt_server = server.PromptServer.instance + + @prompt_server.routes.post("/kiko_local_image_loader/set_node_selection") + async def set_node_selection(request): + """API endpoint to set node selection.""" + try: + data = await request.json() + node_id = str(data.get("node_id")) + path = data.get("path") + media_type = data.get("type") + + if not all([node_id, path, media_type]): + return web.json_response( + {"status": "error", "message": "Missing required data."}, status=400 + ) + + selections = load_selections() + if node_id not in selections: + selections[node_id] = {} + + selections[node_id][media_type] = {"path": path} + save_selections(selections) + + return web.json_response({"status": "ok"}) + except Exception as e: + return web.json_response({"status": "error", "message": str(e)}, status=500) + + @prompt_server.routes.get("/kiko_local_image_loader/get_saved_paths") + async def get_saved_paths(request): + """API endpoint to get saved directory paths.""" + config = load_config() + return web.json_response({"saved_paths": config.get("saved_paths", [])}) + + @prompt_server.routes.post("/kiko_local_image_loader/save_paths") + async def save_paths(request): + """API endpoint to save directory paths.""" + try: + data = await request.json() + paths = data.get("paths", []) + config = load_config() + config["saved_paths"] = paths + save_config(config) + return web.json_response({"status": "ok"}) + except Exception as e: + return web.json_response({"status": "error", "message": str(e)}, status=500) + + @prompt_server.routes.get("/kiko_local_image_loader/images") + async def get_local_images(request): + """API endpoint to get images from a directory.""" + directory = request.query.get("directory", "") + + if not directory or not os.path.isdir(directory): + return web.json_response({"error": "Directory not found."}, status=404) + + # Save last path + config = load_config() + config["last_path"] = directory + save_config(config) + + show_videos = request.query.get("show_videos", "false").lower() == "true" + show_audio = request.query.get("show_audio", "false").lower() == "true" + + page = int(request.query.get("page", 1)) + per_page = int(request.query.get("per_page", 50)) + sort_by = request.query.get("sort_by", "name") + sort_order = request.query.get("sort_order", "asc") + + try: + items = scan_directory( + directory, show_videos, show_audio, sort_by, sort_order + ) + + # Get parent directory + parent_directory = os.path.dirname(directory) + if parent_directory == directory: + parent_directory = None + + # Paginate results + start = (page - 1) * per_page + end = start + per_page + paginated_items = items[start:end] + + return web.json_response( + { + "items": paginated_items, + "total_pages": (len(items) + per_page - 1) // per_page, + "current_page": page, + "current_directory": directory, + "parent_directory": parent_directory, + } + ) + except Exception as e: + return web.json_response({"error": str(e)}, status=500) + + @prompt_server.routes.get("/kiko_local_image_loader/get_last_path") + async def get_last_path(request): + """API endpoint to get last used directory path.""" + return web.json_response({"last_path": load_config().get("last_path", "")}) + + @prompt_server.routes.get("/kiko_local_image_loader/thumbnail") + async def get_thumbnail(request): + """API endpoint to get image thumbnail.""" + filepath = request.query.get("filepath") + if not filepath or ".." in filepath: + return web.Response(status=400) + + filepath = urllib.parse.unquote(filepath) + if not os.path.exists(filepath): + return web.Response(status=404) + + try: + img = Image.open(filepath) + has_alpha = img.mode == "RGBA" or ( + img.mode == "P" and "transparency" in img.info + ) + img = img.convert("RGBA") if has_alpha else img.convert("RGB") + img.thumbnail([320, 320], Image.LANCZOS) + + buffer = io.BytesIO() + format, content_type = ( + ("PNG", "image/png") if has_alpha else ("JPEG", "image/jpeg") + ) + img.save(buffer, format=format, quality=90 if format == "JPEG" else None) + buffer.seek(0) + + return web.Response(body=buffer.read(), content_type=content_type) + except Exception as e: + print(f"KikoLocalImageLoader: Error generating thumbnail: {e}") + return web.Response(status=500) + + @prompt_server.routes.get("/kiko_local_image_loader/view") + async def view_image(request): + """API endpoint to view full image.""" + filepath = request.query.get("filepath") + if not filepath or ".." in filepath: + return web.Response(status=400) + + filepath = urllib.parse.unquote(filepath) + if not os.path.exists(filepath): + return web.Response(status=404) + + try: + return web.FileResponse(filepath) + except Exception: + return web.Response(status=500) + +except ImportError: + # Server not available during testing + pass diff --git a/kikotools/tools/local_image_loader/selections.json b/kikotools/tools/local_image_loader/selections.json new file mode 100644 index 0000000..db6461f --- /dev/null +++ b/kikotools/tools/local_image_loader/selections.json @@ -0,0 +1,12 @@ +{ + "57": { + "image": { + "path": "/home/vito/ai-apps/ComfyUI-3.12/output/2025-05-01/HiDream_00001_.png" + } + }, + "58": { + "image": { + "path": "/home/vito/ai-apps/ComfyUI-3.12/output/CharacterName_00016_.png" + } + } +} \ No newline at end of file