Upload whole batches, cache Drive lookups, handle placeholder credentials

- Only the first image of a batch was uploaded; the rest were silently
  dropped. Every image is now uploaded (single images keep the plain
  "<prefix>.<ext>" name, batches are numbered _001, _002, ...).
- Folder names containing an apostrophe broke the Drive search query.
- The placeholder json/credentials.json and token.json shipped with the
  node produced confusing JSON errors; they're now detected and the node
  says exactly which file to replace.
- Credentials and the Drive client are reused between runs, and folder IDs
  are cached per path, so repeat saves skip the folder lookups (previously
  every save re-read the token, rebuilt the client and searched each folder).
- Images under 5 MB upload in a single request instead of a resumable
  upload session.
- Works with GPU tensors (.cpu() before .numpy()).
- Removed unused imports and the WEB_DIRECTORY that pointed at a
  nonexistent folder; version 1.2.0 in pyproject.toml and __init__.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013Lm4YDXhwj49hAA1RKS58Q

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

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  
 

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 "<prefix>.<ext>" 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)"
-}
+}

__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']

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]