diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7a60b85 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +*.pyc diff --git a/README.md b/README.md index 281e718..7ac68f4 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ The node will create folders you specify in the node itself if they aren't alrea **Folder structure**: Rootfolder/subfolder/subfolder -**Filename prefix**: All images with have the same name, so if the prefix is for example "IMG" all your saved images will be named IMG.png (alternatively IMG.jpeg or IMG.webp) +**Filename prefix**: All images will have the same name, so if the prefix is for example "IMG" all your saved images will be named IMG.png (alternatively IMG.jpeg or IMG.webp). When you send a batch of several images, they are numbered: IMG_001.png, IMG_002.png and so on. **File Format**: Format available is PNG, JPEG and WEBP diff --git a/SaveImageToGoogleDrive.py b/SaveImageToGoogleDrive.py index 852c33f..f3df030 100644 --- a/SaveImageToGoogleDrive.py +++ b/SaveImageToGoogleDrive.py @@ -1,15 +1,10 @@ import os -import comfy.sd -import comfy.utils -import torch import numpy as np from PIL import Image from PIL.PngImagePlugin import PngInfo -import time import threading import json -import random -import string +import uuid import folder_paths from google.auth.transport.requests import Request @@ -19,9 +14,37 @@ from googleapiclient.discovery import build from googleapiclient.http import MediaFileUpload SCOPES = ['https://www.googleapis.com/auth/drive.file'] +DEFAULT_FOLDER = "ComfyUI_Saves" + +# Files under this size are sent in a single request; bigger ones use a resumable upload. +RESUMABLE_THRESHOLD_BYTES = 5 * 1024 * 1024 + +NODE_DIR = os.path.dirname(os.path.abspath(__file__)) +CREDENTIALS_PATH = os.path.join(NODE_DIR, "json", "credentials.json") +TOKEN_PATH = os.path.join(NODE_DIR, "json", "token.json") + + +def _read_json(path): + """Returns the parsed file, or None if it's missing or still the placeholder text shipped with the node.""" + try: + with open(path, "r", encoding="utf-8") as f: + return json.load(f) + except (FileNotFoundError, json.JSONDecodeError): + return None + + +def _escape_query_value(value): + # Drive query strings are quoted with '; a folder name like "Bob's renders" would otherwise break the query. + return value.replace("\\", "\\\\").replace("'", "\\'") class SaveImageToGoogleDrive: + # Shared across runs: the Drive client and the folder IDs already looked up. + _service = None + _creds = None + _folder_cache = {} + _lock = threading.Lock() + def __init__(self): self.output_dir = folder_paths.get_temp_directory() self.type = "temp" @@ -31,7 +54,7 @@ class SaveImageToGoogleDrive: return { "required": { "image": ("IMAGE",), - "google_drive_folder_path": ("STRING", {"default": "ComfyUI_Saves"}), + "google_drive_folder_path": ("STRING", {"default": DEFAULT_FOLDER}), "filename_prefix": ("STRING", {"default": "ComfyUI_Image_"}), "file_format": (["PNG", "JPEG", "WEBP"], {"default": "PNG"}), "quality": ("INT", {"default": 90, "min": 1, "max": 100, "step": 1}), @@ -52,19 +75,17 @@ class SaveImageToGoogleDrive: CATEGORY = "Creepybits/Google Drive" def authenticate_google_drive(self): - node_dir = os.path.dirname(os.path.abspath(__file__)) - credentials_path = os.path.join(node_dir, "json", "credentials.json") - token_path = os.path.join(node_dir, "json", "token.json") + os.makedirs(os.path.dirname(TOKEN_PATH), exist_ok=True) - os.makedirs(os.path.join(node_dir, "json"), exist_ok=True) - - creds = None - if os.path.exists(token_path): - try: - creds = Credentials.from_authorized_user_file(token_path, SCOPES) - except Exception as e: - print(f"Google Drive API Warning: Could not load token file {token_path}. Will attempt re-authentication. Error: {e}") - creds = None + creds = self.__class__._creds + if creds is None: + token_info = _read_json(TOKEN_PATH) + if token_info: + try: + creds = Credentials.from_authorized_user_info(token_info, SCOPES) + except Exception as e: + print(f"Google Drive API Warning: Could not load token file {TOKEN_PATH}. Will attempt re-authentication. Error: {e}") + creds = None if not creds or not creds.valid: if creds and creds.expired and creds.refresh_token: @@ -73,52 +94,74 @@ class SaveImageToGoogleDrive: except Exception as e: print(f"Google Drive API Error refreshing token. Will initiate full re-authentication flow. Error: {e}") creds = None - if not creds: + if not creds or not creds.valid: + client_config = _read_json(CREDENTIALS_PATH) + if not client_config: + print(f"Google Drive API Error: {CREDENTIALS_PATH} is missing or still the placeholder. " + "Download your OAuth client file from Google Cloud Console and save it there.") + return None try: - flow = InstalledAppFlow.from_client_secrets_file(credentials_path, SCOPES) + flow = InstalledAppFlow.from_client_config(client_config, SCOPES) creds = flow.run_local_server(port=0) - except FileNotFoundError: - print(f"Google Drive API Error: credentials.json not found at {credentials_path}. Please make sure you have downloaded it from Google Cloud Console and placed it there.") - return None except Exception as e: - print(f"Google Drive API Error during authentication flow: {e}") - return None + print(f"Google Drive API Error during authentication flow: {e}") + return None try: - with open(token_path, 'w') as token: + with open(TOKEN_PATH, 'w') as token: token.write(creds.to_json()) except Exception as e: - print(f"Google Drive API Warning: Could not save token file to {token_path}. You may need to re-authenticate next time. Error: {e}") + print(f"Google Drive API Warning: Could not save token file to {TOKEN_PATH}. You may need to re-authenticate next time. Error: {e}") + # New credentials need a new client. + self.__class__._service = None + + self.__class__._creds = creds return creds - def save_image(self, image, google_drive_folder_path, filename_prefix, file_format, quality, prompt=None, extra_pnginfo=None): - status_message = "Google Drive Save: Starting..." - print(status_message) + def get_service(self, creds): + if self.__class__._service is None: + self.__class__._service = build('drive', 'v3', credentials=creds, cache_discovery=False) + return self.__class__._service - if image.shape[0] > 1: - image_to_process = image[0].unsqueeze(0) - else: - image_to_process = image - image_np = image_to_process[0].numpy() * 255.0 - image_pil = Image.fromarray(np.clip(image_np, 0, 255).astype(np.uint8)) + def get_folder_id(self, service, path_parts): + """Finds (or creates) each folder in the path. Results are cached so repeated saves to + the same folder skip the lookups.""" + cache_key = "/".join(path_parts) + cached = self._folder_cache.get(cache_key) + if cached: + return cached + + current_parent_id = 'root' + for folder_name in path_parts: + folder_exists_query = (f"name='{_escape_query_value(folder_name)}' and mimeType='application/vnd.google-apps.folder' " + f"and '{current_parent_id}' in parents and trashed=false") + response = service.files().list(q=folder_exists_query, + spaces='drive', + fields='files(id, name)', + pageSize=1).execute() + found_folders = response.get('files', []) + + if found_folders: + current_parent_id = found_folders[0]['id'] + else: + folder_metadata = { + 'name': folder_name, + 'mimeType': 'application/vnd.google-apps.folder', + 'parents': [current_parent_id] + } + created_folder = service.files().create(body=folder_metadata, fields='id').execute() + current_parent_id = created_folder.get('id') + + self._folder_cache[cache_key] = current_parent_id + return current_parent_id + + def save_image(self, image, google_drive_folder_path, filename_prefix, file_format, quality, prompt=None, extra_pnginfo=None): + print("Google Drive Save: Starting...") chosen_format_upper = file_format.upper() - file_extension = "png" - - if chosen_format_upper == "JPEG": - file_extension = "jpeg" - if image_pil.mode == 'RGBA': - image_pil = image_pil.convert("RGB") - elif chosen_format_upper == "WEBP": - file_extension = "webp" - if image_pil.mode == 'P': - image_pil = image_pil.convert("RGBA") - elif image_pil.mode == 'RGBA' and not image_pil.getchannel('A').getbbox(): - image_pil = image_pil.convert("RGB") - - temp_filename_unique = f"comfy_gdrive_temp_{os.getpid()}_{threading.current_thread().ident}_{int(time.time())}_{''.join(random.choice(string.ascii_lowercase + string.digits) for i in range(8))}.{file_extension}" - temp_filepath_local = os.path.join(self.output_dir, temp_filename_unique) + file_extension = {"JPEG": "jpeg", "WEBP": "webp"}.get(chosen_format_upper, "png") + upload_mimetype = {"PNG": "image/png", "JPEG": "image/jpeg", "WEBP": "image/webp"}[chosen_format_upper] metadata = None if chosen_format_upper == "PNG": @@ -129,112 +172,71 @@ class SaveImageToGoogleDrive: for k, v in extra_pnginfo.items(): metadata.add_text(k, json.dumps(v)) - try: - os.makedirs(self.output_dir, exist_ok=True) - if chosen_format_upper == "PNG": - image_pil.save(temp_filepath_local, pnginfo=metadata, compress_level=4) - elif chosen_format_upper in ["JPEG", "WEBP"]: - image_pil.save(temp_filepath_local, quality=quality) - else: - image_pil.save(temp_filepath_local) - image_pil.close() - except Exception as e: - status_message = f"Google Drive Save Error: Could not save temporary image file to {temp_filepath_local}. {e}" - print(status_message) - return {"ui": {"images": []}, "result": (image, status_message,)} + # Save every image in the batch locally (also used for the preview). This used to upload only the first one. + os.makedirs(self.output_dir, exist_ok=True) + local_files = [] + for i in range(image.shape[0]): + image_pil = Image.fromarray(np.clip(image[i].cpu().numpy() * 255.0, 0, 255).astype(np.uint8)) + if chosen_format_upper == "JPEG" and image_pil.mode != "RGB": + image_pil = image_pil.convert("RGB") - preview_dict = { - "filename": temp_filename_unique, - "subfolder": "", - "type": self.type - } - - creds = self.authenticate_google_drive() - if not creds: - status_message = "Google Drive Save Error: Authentication failed. Check console for details." - print(status_message) - return {"ui": {"images": [preview_dict]}, "result": (image, status_message,)} - - target_folder_id = None - try: - service = build('drive', 'v3', credentials=creds) - - if not google_drive_folder_path or google_drive_folder_path.strip() == "" or google_drive_folder_path.strip() == "/": - path_parts = [self.INPUT_TYPES()["required"]["google_drive_folder_path"][1]["default"]] - else: - path_parts = [part for part in google_drive_folder_path.strip('/').split('/') if part] - if not path_parts: - path_parts = [self.INPUT_TYPES()["required"]["google_drive_folder_path"][1]["default"]] - - current_parent_id = 'root' - - for i, folder_name in enumerate(path_parts): - folder_exists_query = f"name='{folder_name}' and mimeType='application/vnd.google-apps.folder' and '{current_parent_id}' in parents and trashed=false" - response = service.files().list(q=folder_exists_query, - spaces='drive', - fields='files(id, name)', - pageSize=1).execute() - found_folders = response.get('files', []) - - if found_folders: - current_folder_id = found_folders[0]['id'] + temp_filename_unique = f"comfy_gdrive_temp_{uuid.uuid4().hex}.{file_extension}" + temp_filepath_local = os.path.join(self.output_dir, temp_filename_unique) + try: + if chosen_format_upper == "PNG": + image_pil.save(temp_filepath_local, pnginfo=metadata, compress_level=4) else: - folder_metadata = { - 'name': folder_name, - 'mimeType': 'application/vnd.google-apps.folder', - 'parents': [current_parent_id] - } - created_folder = service.files().create(body=folder_metadata, fields='id').execute() - current_folder_id = created_folder.get('id') - - current_parent_id = current_folder_id - if i == len(path_parts) - 1: - target_folder_id = current_folder_id - - if not target_folder_id: - status_message = f"Google Drive Save Error: Could not determine target folder ID for path '{google_drive_folder_path}'." + image_pil.save(temp_filepath_local, quality=quality) + except Exception as e: + status_message = f"Google Drive Save Error: Could not save temporary image file to {temp_filepath_local}. {e}" print(status_message) - return {"ui": {"images": [preview_dict]}, "result": (image, status_message,)} + return {"ui": {"images": []}, "result": (image, status_message,)} + local_files.append(temp_filename_unique) - except Exception as e: - status_message = f"Google Drive Save Error processing folder path '{google_drive_folder_path}'. {e}" - print(status_message) - return {"ui": {"images": [preview_dict]}, "result": (image, status_message,)} + previews = [{"filename": f, "subfolder": "", "type": self.type} for f in local_files] + path_parts = [part for part in google_drive_folder_path.strip().strip('/').split('/') if part.strip()] or [DEFAULT_FOLDER] + display_path = "/".join(path_parts) - uploaded_file_id = None - try: - final_filename = f"{filename_prefix}.{file_extension}" - file_metadata = { - 'name': final_filename, - 'parents': [target_folder_id] - } + with self._lock: + creds = self.authenticate_google_drive() + if not creds: + status_message = "Google Drive Save Error: Authentication failed. Check console for details." + print(status_message) + return {"ui": {"images": previews}, "result": (image, status_message,)} - mimetype_map = { - "PNG": "image/png", - "JPEG": "image/jpeg", - "WEBP": "image/webp" - } - upload_mimetype = mimetype_map.get(chosen_format_upper, 'application/octet-stream') + try: + service = self.get_service(creds) + target_folder_id = self.get_folder_id(service, path_parts) + except Exception as e: + # A cached folder may have been deleted on Drive; look it up fresh next time. + self._folder_cache.pop(display_path, None) + status_message = f"Google Drive Save Error processing folder path '{google_drive_folder_path}'. {e}" + print(status_message) + return {"ui": {"images": previews}, "result": (image, status_message,)} - media = MediaFileUpload(temp_filepath_local, - mimetype=upload_mimetype, - resumable=True, - chunksize=1024*1024) - file = service.files().create(body=file_metadata, - media_body=media, - fields='id,webContentLink,webViewLink').execute() - uploaded_file_id = file.get('id') + statuses = [] + for i, temp_filename in enumerate(local_files): + # A single image keeps the plain "." name; batches are numbered. + final_filename = f"{filename_prefix}.{file_extension}" if len(local_files) == 1 else f"{filename_prefix}_{i + 1:03d}.{file_extension}" + temp_filepath_local = os.path.join(self.output_dir, temp_filename) + try: + media = MediaFileUpload(temp_filepath_local, + mimetype=upload_mimetype, + resumable=os.path.getsize(temp_filepath_local) > RESUMABLE_THRESHOLD_BYTES, + chunksize=1024*1024) + file = service.files().create(body={'name': final_filename, 'parents': [target_folder_id]}, + media_body=media, + fields='id').execute() + statuses.append(f"Successfully uploaded '{final_filename}' to folder '{display_path}'. File ID: {file.get('id')}") + except Exception as e: + if "notFound" in str(e) or "File not found" in str(e): + self._folder_cache.pop(display_path, None) + statuses.append(f"Error: Could not upload file '{final_filename}'. {e}") + print(f"Google Drive Save: {statuses[-1]}") - display_path = "/".join(path_parts) - status_message = f"Google Drive Save: Successfully uploaded '{final_filename}' to folder '{display_path}'. File ID: {uploaded_file_id}" - print(status_message) - - except Exception as e: - status_message = f"Google Drive Save Error: Could not upload file '{os.path.basename(temp_filepath_local)}'. {e}" - print(status_message) - - return {"ui": {"images": [preview_dict]}, "result": (image, status_message,)} + status_message = "Google Drive Save: " + " | ".join(statuses) + return {"ui": {"images": previews}, "result": (image, status_message,)} NODE_CLASS_MAPPINGS = { @@ -243,4 +245,4 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "SaveImageToGoogleDrive": "Save Image to Google Drive (Creepybits)" -} \ No newline at end of file +} diff --git a/__init__.py b/__init__.py index 823f8de..058a80a 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,6 @@ -import os -import sys from .SaveImageToGoogleDrive import NODE_CLASS_MAPPINGS as nodes_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as nodes_NODE_DISPLAY_NAME_MAPPINGS -__version__ = "1.1.1" - -# Define the web directory for ComfyUI to find our JavaScript files -WEB_DIRECTORY = "./web" +__version__ = "1.2.0" NODE_CLASS_MAPPINGS = { **nodes_NODE_CLASS_MAPPINGS, @@ -15,4 +10,4 @@ NODE_DISPLAY_NAME_MAPPINGS = { **nodes_NODE_DISPLAY_NAME_MAPPINGS, } -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/pyproject.toml b/pyproject.toml index 46477f4..5fe0203 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "Save_To_GDrive" description = "A ComfyUI node that lets you save your images and videos directly to Google Drive through Google's free API service." -version = "1.0.0" +version = "1.2.0" license = {text = "MIT License"} [project.urls]