diff --git a/README.md b/README.md index 4a6321b..fc03db7 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,8 @@ workflows, especially when dealing with multiple audio inputs or outputs. - [15. Audio Normalize (Peak)](#15-audio-normalize-peak) - [16. Audio Apply Batched Gain](#16-audio-apply-batched-gain) - [17. Audio Download and Load](#17-audio-download-and-load) +- [✨ Extra nodes](#-extra-nodes) + - [1. Image Download and Load](#1-image-download-and-load) - [🚀 Installation](#-installation) - [📦 Dependencies](#-dependencies) - [🖼️ Examples](#️-examples) @@ -303,6 +305,28 @@ workflows, especially when dealing with multiple audio inputs or outputs. - **Player**: This node uses an `AUDIO_UI` widget to allow playing the downloaded song. +## ✨ Extra nodes + +These nodes are just complements for the other nodes. I.e. they share a lot of code. + +### 1. Image Download and Load + - **Display Name:** `Image Download and Load` + - **Internal Name:** `SET_ImageDownload` + - **Category:** `image/io` + - **Description:** Downloads an image file from a URL into the `ComfyUI/input/` directory if it's not already there, and then loads it as an image and mask. This is perfect for creating self-contained, shareable workflows with example image. + - **Inputs:** + - `image_bypass` (IMAGE, Optional): If an image is provided here it will be used for the output. You can connect a `Load Image` node here, if the connected node is muted (bypassed) we download the file, otherwise we use the image from the `Load Image` node. + - `mask_bypass` (MASK, Optional): This input complements `image_bypass`. + - `base_url` (STRING): The URL of the directory containing the image file. + - `filename` (STRING): The name of the image file to download (e.g., photo.jpg, art.png). + - **Output:** + - `image` (IMAGE): The loaded image. + - `alpha_mask` (MASK): The alpha mask for the loaded image. + - **Behavior Details:** + - **Caching:** The node checks the `ComfyUI/input/` folder first. If the file with the specified `filename` already exists, the download is skipped. + - **Bypass:** If only one of `image_bypass` and `mask_bypass` is connected the other will be assumed to be empty. You should connect both or avoid using the output corresponding to the unconnected input. + + ## 🚀 Installation You can install the nodes from the ComfyUI nodes manager, the name is *Audio Batch*, or just do it manually: diff --git a/__init__.py b/__init__.py index b061516..2ec390c 100644 --- a/__init__.py +++ b/__init__.py @@ -14,7 +14,8 @@ NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} for name, obj in inspect.getmembers(nodes_audio): - if not inspect.isclass(obj) or not hasattr(obj, "INPUT_TYPES"): + # We skip nodes imported from the ComfyUI main nodes + if not inspect.isclass(obj) or not hasattr(obj, "INPUT_TYPES") or obj.__module__ == "nodes": continue assert hasattr(obj, "UNIQUE_NAME"), f"No name for {obj.__name__}" NODE_CLASS_MAPPINGS[obj.UNIQUE_NAME] = obj diff --git a/source/nodes/nodes_audio.py b/source/nodes/nodes_audio.py index 2ef3f9f..464f2b3 100644 --- a/source/nodes/nodes_audio.py +++ b/source/nodes/nodes_audio.py @@ -21,6 +21,12 @@ except ModuleNotFoundError: # No ComfyUI, this is a test environment def get_input_directory(): return "" +try: + # We need to import the built-in LoadImage class for ImageDownload + from nodes import LoadImage # type: ignore[attr-defined] + has_load_image = True +except Exception: + has_load_image = False logger = main_logger BASE_CATEGORY = "audio" @@ -1096,7 +1102,7 @@ class AudioDownload: send_toast_notification(f"Downloading `{filename}`", "Download") try: - download_model(url=download_url, save_dir=save_dir, file_name=filename) + download_model(url=download_url, save_dir=save_dir, file_name=filename, kind="audio") except Exception as e: logger.error(f"Download failed for {download_url}: {e}", exc_info=True) raise # Re-raise to stop the workflow and show the error @@ -1143,3 +1149,100 @@ class AudioDownload: except Exception as e: logger.error(f"Failed to load or process audio file '{local_filepath}': {e}", exc_info=True) raise IOError(f"Could not load or process the audio file: {filename}. It may be missing or corrupt.") from e + + +if has_load_image: + class ImageDownload: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "base_url": ("STRING", { + "default": + "https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/", + "tooltip": "The base URL where the image file is located." + }), + "filename": ("STRING", { + "default": "audioseparation_logo.jpg", + "tooltip": "The name of the image file to download (e.g., photo.jpg, art.png)." + }), + }, + "optional": { + "image_bypass": ("IMAGE", { + "tooltip": "If this image is present will be used instead of the downloaded one" + }), + "mask_bypass": ("MASK", {"tooltip": "If this mask is present will be used instead of the downloaded one"}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + RETURN_NAMES = ("image", "alpha_mask") + FUNCTION = "load_or_download_image" + CATEGORY = "image" + "/" + IO_CATEGORY + DESCRIPTION = ("Downloads an image to ComfyUI's 'input' directory if it doesn't exist, then loads it using the " + "built-in LoadImage logic.") + UNIQUE_NAME = "SET_ImageDownload" + DISPLAY_NAME = "Image Download and Load" + OUTPUT_NODE = True + + def load_or_download_image(self, base_url: str, filename: str, image_bypass: Optional[torch.Tensor] = None, + mask_bypass: Optional[torch.Tensor] = None): + # If we have something at the bypass inputs use it + if image_bypass is not None or mask_bypass is not None: + if image_bypass is None: + # Just a mask + assert mask_bypass is not None, "This should not be possible if image_bypass is None" # For mypy + image_bypass = torch.zeros(mask_bypass.shape + (3,), dtype=torch.float32, device="cpu") + logger.warning("ImageDownload: Returning an empty image") + elif mask_bypass is None: + # This is ComfyUI behavior when we don't have transparency + mask_bypass = torch.zeros((64, 64), dtype=torch.float32, device="cpu").unsqueeze(0) + logger.warning("ImageDownload: Returning an empty mask") + return (image_bypass, mask_bypass) + + save_dir = get_input_directory() + local_filepath = os.path.join(save_dir, filename) + + if not os.path.exists(local_filepath): + logger.info(f"File '{filename}' not found locally. Attempting to download.") + + if not base_url.endswith('/'): + base_url += '/' + download_url = base_url + filename + + try: + download_model(url=download_url, save_dir=save_dir, file_name=filename, kind="image") + except Exception as e: + logger.error(f"Download failed for {download_url}: {e}", exc_info=True) + raise + + else: + logger.info(f"Found existing file, skipping download: '{local_filepath}'") + + # --- REUSE ComfyUI's LoadImage LOGIC --- + try: + # Instantiate the built-in LoadImage node + loader_instance = LoadImage() + + # The LoadImage node's `load_image` method expects the filename as passed + # by the ComfyUI widget, which is just the filename. It internally + # resolves the path using folder_paths. + + logger.debug(f"Calling built-in LoadImage.load_image() with filename: '{filename}'") + + # Call the method and return its result directly + result = loader_instance.load_image(filename) + downloaded_file = { + "images": [{ + "filename": filename, + "subfolder": "", + "type": "input" # We stored the file in the "input" folder + }] + } + return {"ui": downloaded_file, "result": result} + + except Exception as e: + logger.error(f"Failed to load image '{filename}' using built-in LoadImage node: {e}", exc_info=True) + # Re-raise to make the error visible in ComfyUI + raise IOError(f"Could not load the image file '{filename}' using the standard loader. " + "It may be corrupt or in an unsupported format.") from e diff --git a/source/nodes/utils/downloader.py b/source/nodes/utils/downloader.py index 430f491..c47d2de 100644 --- a/source/nodes/utils/downloader.py +++ b/source/nodes/utils/downloader.py @@ -191,8 +191,8 @@ def download_model_urllib(url: str, save_dir: str, file_name: str): return Downloader(save_dir, file_name).download_model(url) -def download_model(url: str, save_dir: str, file_name: str, force_urllib: bool = False): - logger.info(f"Downloading model: {file_name}") +def download_model(url: str, save_dir: str, file_name: str, force_urllib: bool = False, kind: str = "model"): + logger.info(f"Downloading {kind}: {file_name}") logger.info(f"Source URL: {url}") full_name = os.path.join(save_dir, file_name) logger.info(f"Destination: {full_name}")