diff --git a/cache/1zNxWo6TrxC73_uosnjVoHBl7q6u0eMy0.png b/cache/1zNxWo6TrxC73_uosnjVoHBl7q6u0eMy0.png new file mode 100644 index 0000000..c6f3404 Binary files /dev/null and b/cache/1zNxWo6TrxC73_uosnjVoHBl7q6u0eMy0.png differ diff --git a/cache/image_cache_index.json b/cache/image_cache_index.json new file mode 100644 index 0000000..efa729a --- /dev/null +++ b/cache/image_cache_index.json @@ -0,0 +1 @@ +{"1zNxWo6TrxC73_uosnjVoHBl7q6u0eMy0": "/Users/fillipisgro/Desktop/Git Projects/ComfyUI/custom_nodes/ComfyUI_Fill-Nodes/cache/1zNxWo6TrxC73_uosnjVoHBl7q6u0eMy0.png"} \ No newline at end of file diff --git a/nodes/FL_GoogleDriveImageDownloader.py b/nodes/FL_GoogleDriveImageDownloader.py index 526f3d9..f1d61f3 100644 --- a/nodes/FL_GoogleDriveImageDownloader.py +++ b/nodes/FL_GoogleDriveImageDownloader.py @@ -4,6 +4,12 @@ import gdown from PIL import Image import torch import numpy as np +import hashlib +import json +from pathlib import Path +import hashlib +import json +from pathlib import Path class FL_GoogleDriveImageDownloader: @@ -16,6 +22,10 @@ class FL_GoogleDriveImageDownloader: "multiline": False, "placeholder": "Enter Google Drive image link" }), + "use_cache": ("BOOLEAN", { + "default": True, + "label": "Use Cached Image" + }), }, } @@ -23,6 +33,16 @@ class FL_GoogleDriveImageDownloader: FUNCTION = "download_and_process_image" CATEGORY = "🏵️Fill Nodes/Google Drive" + def __init__(self): + self.cache_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "cache") + self.cache_index_file = os.path.join(self.cache_dir, "image_cache_index.json") + os.makedirs(self.cache_dir, exist_ok=True) + + # Initialize the cache index if it doesn't exist + if not os.path.exists(self.cache_index_file): + with open(self.cache_index_file, "w") as f: + json.dump({}, f) + def extract_file_id_from_link(self, share_link: str) -> str: """Extracts the file ID from a Google Drive share link.""" if not share_link: @@ -35,48 +55,121 @@ class FL_GoogleDriveImageDownloader: return match.group(1) raise ValueError("Unable to extract file ID. Please ensure you're using a valid sharing link") - def download_and_process_image(self, google_drive_link: str) -> tuple: + def get_cached_image(self, file_id): + """Check if image is in cache and return the path if it exists""" try: - # Create temporary directory for download - temp_dir = os.path.join(os.getcwd(), 'temp_downloads') - os.makedirs(temp_dir, exist_ok=True) + with open(self.cache_index_file, "r") as f: + cache_index = json.load(f) + + if file_id in cache_index: + cached_path = cache_index[file_id] + if os.path.exists(cached_path): + return cached_path + except Exception as e: + print(f"Cache lookup error: {str(e)}") + + return None + + def save_to_cache(self, file_id, image_path): + """Save downloaded image to cache""" + try: + # Generate a unique filename + cache_filename = os.path.join(self.cache_dir, f"{file_id}.png") + + # Copy the image to cache + img = Image.open(image_path) + img.save(cache_filename) + + # Update the cache index + with open(self.cache_index_file, "r") as f: + cache_index = json.load(f) + + cache_index[file_id] = cache_filename + + with open(self.cache_index_file, "w") as f: + json.dump(cache_index, f) + + return cache_filename + except Exception as e: + print(f"Cache save error: {str(e)}") + return None - # Extract file ID and create download URL + def download_and_process_image(self, google_drive_link: str, use_cache: bool = True) -> tuple: + try: + # Extract file ID file_id = self.extract_file_id_from_link(google_drive_link) - url = f'https://drive.google.com/uc?id={file_id}' - - # Download the file - print("Downloading image from Google Drive...") - temp_path = os.path.join(temp_dir, f"temp_image_{file_id}") - output = gdown.download(url=url, output=temp_path, quiet=False, fuzzy=True) - - if not output: - raise ValueError("Failed to download image") - - # Open and verify it's an image + + image_path = None + temp_path = None # Initialize temp_path to ensure it's always defined + + # Check cache first if caching is enabled + if use_cache: + cached_path = self.get_cached_image(file_id) + if cached_path: + print(f"Using cached image for file ID: {file_id}") + image_path = cached_path + + # If not in cache or caching disabled, download the image + if not image_path: + # Create temporary directory for download + temp_dir = os.path.join(os.getcwd(), 'temp_downloads') + os.makedirs(temp_dir, exist_ok=True) + + # Create download URL + url = f'https://drive.google.com/uc?id={file_id}' + + # Download the file + print("Downloading image from Google Drive...") + temp_path = os.path.join(temp_dir, f"temp_image_{file_id}") + output = gdown.download(url=url, output=temp_path, quiet=False, fuzzy=True) + + if not output: + raise ValueError("Failed to download image") + + image_path = output + + # Save to cache if enabled + if use_cache: + cached_path = self.save_to_cache(file_id, image_path) + if cached_path: + image_path = cached_path + + # Open and process the image try: - image = Image.open(output) + image = Image.open(image_path) # Convert to RGB if necessary if image.mode != 'RGB': image = image.convert('RGB') except Exception as e: - raise ValueError(f"Downloaded file is not a valid image: {str(e)}") - + raise ValueError(f"Invalid image file: {str(e)}") + # Convert to the format expected by ComfyUI - # ComfyUI expects images as: torch.Tensor [B, H, W, C] in float32 image_np = np.array(image).astype(np.float32) / 255.0 image_tensor = torch.from_numpy(image_np) image_tensor = image_tensor.unsqueeze(0) # Add batch dimension - - # Cleanup - os.remove(output) - + + # Cleanup temp file if it exists and isn't the cached version + if temp_path and os.path.exists(temp_path) and temp_path != image_path: + os.remove(temp_path) + print(f"Successfully processed image: {image_tensor.shape}") return (image_tensor,) - + except Exception as e: raise ValueError(f"Error processing image from Google Drive: {str(e)}") @classmethod - def IS_CHANGED(cls, google_drive_link): - return float("NaN") \ No newline at end of file + def IS_CHANGED(cls, google_drive_link, use_cache): + # Only signal a change if use_cache is False + # This ensures the node won't rerun when caching is enabled + if not use_cache: + return float("NaN") + + # Otherwise, we should check if the link has changed since last run + # This is done by hashing the link + if not google_drive_link: + return 0 + + # Simple hash of the link string + link_hash = hashlib.md5(google_drive_link.encode()).hexdigest() + return link_hash \ No newline at end of file