diff --git a/example/Basic_Scraper.json b/example/Basic_Scraper.json new file mode 100644 index 0000000..00b8e0e --- /dev/null +++ b/example/Basic_Scraper.json @@ -0,0 +1,166 @@ +{ + "last_node_id": 8, + "last_link_id": 5, + "nodes": [ + { + "id": 1, + "type": "polymath_scraper", + "pos": [ + 130, + 240 + ], + "size": [ + 330, + 216 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 1 + ], + "slot_index": 0, + "shape": 6 + }, + { + "name": "filepath_texts", + "type": "STRING", + "links": [ + 2 + ], + "slot_index": 1, + "shape": 6 + } + ], + "properties": { + "Node name for S&R": "polymath_scraper" + }, + "widgets_values": [ + "https://de.pinterest.com/search/pins/?q=%23cars&rs=typed", + "output/cars", + "cars", + 4, + false, + false + ], + "color": "#9968f3", + "bgcolor": "#519b36" + }, + { + "id": 3, + "type": "ShowText|pysssss", + "pos": [ + 520, + 340 + ], + "size": [ + 450, + 226 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "text", + "type": "STRING", + "link": 2, + "widget": { + "name": "text" + } + } + ], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [], + "slot_index": 0, + "shape": 6 + } + ], + "properties": { + "Node name for S&R": "ShowText|pysssss" + }, + "widgets_values": [ + "", + "/Users/benjaminbertram/ComfyUI/output/cars/cars_00002.png", + "/Users/benjaminbertram/ComfyUI/output/cars/cars_00003.png", + "/Users/benjaminbertram/ComfyUI/output/cars/cars_00004.png", + "/Users/benjaminbertram/ComfyUI/output/cars/cars_00005.png" + ] + }, + { + "id": 2, + "type": "PreviewImage", + "pos": [ + 1020, + 210 + ], + "size": [ + 310, + 330 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 1 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [] + } + ], + "links": [ + [ + 1, + 1, + 0, + 2, + 0, + "IMAGE" + ], + [ + 2, + 1, + 1, + 3, + 0, + "STRING" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 1, + "offset": { + "0": 0, + "1": 0 + } + }, + "node_versions": { + "comfyui_LLM_Polymath": "3e224dd0a34628176137e2a6d24dc57471fc9570", + "ComfyUI-Custom-Scripts": "a53ef9b617ed1331640d7a2cd97644995908dc00", + "comfy-core": "0.3.14" + }, + "ue_links": [], + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/polymath.py b/polymath.py index ae6a295..878c8bb 100644 --- a/polymath.py +++ b/polymath.py @@ -1,7 +1,7 @@ import subprocess import sys import json -import os +import os, glob import re import base64 import numpy as np @@ -14,6 +14,8 @@ from bs4 import BeautifulSoup # For the Scraper import torch +import folder_paths +import shutil api_key = os.getenv("OPENAI_API_KEY") @@ -395,6 +397,7 @@ class Polymath: def pil2tensor(image): return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + class MediaScraper: @classmethod def INPUT_TYPES(s): @@ -405,75 +408,159 @@ class MediaScraper: "file_name": ("STRING", {"multiline": False, "default": ""}), }, "optional": { - "image_load_cap": ("INT", {"default": 0, "min": 0, "step": 1}) + "image_load_cap": ("INT", {"default": 0, "min": 0, "step": 1}), + #"preview_images": ("BOOLEAN", {"default": False}), + "keep_temp_path": ("BOOLEAN", {"default": False, "tooltip": "Images are saved temporarily, then renamed. Set true to keep the gallery_dl structure."}), } } - + RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("images", "filepath_texts") FUNCTION = "execute" CATEGORY = "Polymath" OUTPUT_IS_LIST = (True, True) - def execute(self, urls, output_file_path, file_name, image_load_cap=0): + def save_images_custom(self, images, dest_folder, file_prefix, compress_level=4): + """ + Saves the given tensor images into dest_folder using file_prefix as the name prefix. + Avoids overwriting existing files by incrementing an index. + """ + import folder_paths + import glob + + # If no dest_folder is given, use ComfyUI's default + if not dest_folder: + dest_folder = folder_paths.get_output_directory() + if not os.path.exists(dest_folder): + os.makedirs(dest_folder) + + # Attempt to figure out the (height, width) from the first image + try: + tensor = images[0] + if len(tensor.shape) == 4: + height, width = tensor.shape[1], tensor.shape[2] + elif len(tensor.shape) == 3: + height, width = tensor.shape[0], tensor.shape[1] + else: + height, width = 100, 100 + except: + height, width = 100, 100 + + # Use folder_paths logic for naming + full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path( + file_prefix or "ComfyUI", dest_folder, width, height + ) + existing_files = glob.glob(os.path.join(dest_folder, "*")) + counter += len(existing_files) + + saved_paths = [] + for idx, image_tensor in enumerate(images, start=1): + try: + if isinstance(image_tensor, torch.Tensor): + arr = image_tensor.squeeze(0).cpu().numpy() if image_tensor.ndim == 4 else image_tensor.cpu().numpy() + arr = (255. * arr).clip(0, 255).astype("uint8") + pil_img = Image.fromarray(arr) + else: + pil_img = image_tensor + except Exception: + continue + + file_name = f"{filename}_{counter:05}.png" + out_path = os.path.join(dest_folder, file_name) + # If file already exists, append idx to avoid overwriting + if os.path.exists(out_path): + file_name = f"{filename}_{counter:05}_{idx:02}.png" + out_path = os.path.join(dest_folder, file_name) + pil_img.save(out_path, compress_level=compress_level) + saved_paths.append(out_path) + counter += 1 + + return saved_paths + + def execute(self, urls, output_file_path, file_name, image_load_cap=0, preview_images=False, keep_temp_path=False): + # 1) Create a temp folder for downloads if both output path and file name are given. + if output_file_path and file_name: + temp_folder = os.path.join(os.path.abspath(output_file_path), "__temp__") + os.makedirs(temp_folder, exist_ok=True) + download_dest = temp_folder + else: + # If user didn't specify both, download directly to output_file_path (or current dir if empty) + download_dest = os.path.abspath(output_file_path) if output_file_path else None + + # 2) Download images using gallery-dl url_list = [u.strip() for u in urls.split("\n") if u.strip()] downloaded_files = [] - - # Download step: Run gallery-dl for each URL for url in url_list: if image_load_cap > 0 and len(downloaded_files) >= image_load_cap: break cmd = ["gallery-dl", url] - if output_file_path: - cmd.extend(["-d", output_file_path]) + if download_dest: + cmd.extend(["-d", download_dest]) if image_load_cap > 0: cap = image_load_cap - len(downloaded_files) cmd.extend(["--range", f"1-{cap}"]) subprocess.run(cmd) - # List all downloaded files from the output directory - valid_extensions = ('.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif') - if output_file_path and os.path.isdir(output_file_path): - for root, dirs, files in os.walk(output_file_path): + # 3) Gather downloaded files + valid_exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp', '.gif') + if download_dest and os.path.isdir(download_dest): + for root, dirs, files in os.walk(download_dest): for f in sorted(files): - if f.lower().endswith(valid_extensions): - full_path = os.path.join(root, f) - downloaded_files.append(full_path) + if f.lower().endswith(valid_exts): + downloaded_files.append(os.path.join(root, f)) if image_load_cap > 0 and len(downloaded_files) >= image_load_cap: break if image_load_cap > 0 and len(downloaded_files) >= image_load_cap: break - # Batch rename if file_name is provided - if file_name and downloaded_files: - base, ext = os.path.splitext(file_name) - ext = ext if ext else ".jpg" - for idx, old_path in enumerate(downloaded_files, start=1): - new_name = f"{base}_{idx}{ext}" - new_path = os.path.join(os.path.dirname(old_path), new_name) - os.rename(old_path, new_path) - downloaded_files[idx-1] = new_path - - valid_files = [f for f in downloaded_files if os.path.isfile(f) and f.lower().endswith(valid_extensions)] - - # Initialize empty lists for images and filepath_texts - images = [] - filepath_texts = [] - - for file_path in valid_files: + # 4) Convert downloaded images to tensors + images, filepath_texts = [], [] + for fp in downloaded_files: try: - img = Image.open(file_path) + img = Image.open(fp) images.append(pil2tensor(img)) - filepath_texts.append(file_path) # Add file path to the list + filepath_texts.append(fp) except Exception: pass - if not images: - images.append(torch.ones(3, 100, 100)) - filepath_texts.append("") # Add an empty string if no images are found - + # 5) If user specified both output path + file name + if output_file_path and file_name: + final_folder = os.path.abspath(output_file_path) # We'll store final images here + prefix = os.path.splitext(file_name)[0] + + if keep_temp_path: + # Simply rename in place, keep the temp folder + base, ext = os.path.splitext(file_name) + ext = ext if ext else ".jpg" + for idx, old_file in enumerate(downloaded_files, start=1): + new_file_name = f"{base}_{idx}{ext}" + new_file_path = os.path.join(os.path.dirname(old_file), new_file_name) + os.rename(old_file, new_file_path) + downloaded_files[idx - 1] = new_file_path + filepath_texts = downloaded_files + + else: + # Move (save) images into final_folder with file_name as prefix + # Then remove the temp folder entirely + filepath_texts = self.save_images_custom(images, final_folder, prefix) + for old_file in downloaded_files: + try: + os.remove(old_file) + except: + pass + # Remove temp folder if empty (or remove it entirely) + try: + shutil.rmtree(temp_folder) + except: + pass + + else: + # If not both set, we just keep the downloaded files in place + filepath_texts = downloaded_files + return (images, filepath_texts) + NODE_CLASS_MAPPINGS = { "polymath_chat": Polymath, "polymath_scraper": MediaScraper