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 folder_paths from google.auth.transport.requests import Request from google.oauth2.credentials import Credentials from google_auth_oauthlib.flow import InstalledAppFlow from googleapiclient.discovery import build from googleapiclient.http import MediaFileUpload SCOPES = ['https://www.googleapis.com/auth/drive.file'] class SaveImageToGoogleDrive: def __init__(self): self.output_dir = folder_paths.get_temp_directory() self.type = "temp" @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "google_drive_folder_path": ("STRING", {"default": "ComfyUI_Saves"}), "filename_prefix": ("STRING", {"default": "ComfyUI_Image_"}), "file_format": (["PNG", "JPEG", "WEBP"], {"default": "PNG"}), "quality": ("INT", {"default": 90, "min": 1, "max": 100, "step": 1}), }, "hidden": { "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" }, } RETURN_TYPES = ("IMAGE", "STRING",) RETURN_NAMES = ("image", "status",) FUNCTION = "save_image" OUTPUT_NODE = True 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.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 if not creds or not creds.valid: if creds and creds.expired and creds.refresh_token: try: creds.refresh(Request()) 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: try: flow = InstalledAppFlow.from_client_secrets_file(credentials_path, 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 try: 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}") 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) 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)) 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) metadata = None if chosen_format_upper == "PNG": metadata = PngInfo() if prompt is not None: metadata.add_text("prompt", json.dumps(prompt)) if extra_pnginfo is not None: 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,)} 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'] 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}'." print(status_message) return {"ui": {"images": [preview_dict]}, "result": (image, status_message,)} 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,)} uploaded_file_id = None try: final_filename = f"{filename_prefix}.{file_extension}" file_metadata = { 'name': final_filename, 'parents': [target_folder_id] } mimetype_map = { "PNG": "image/png", "JPEG": "image/jpeg", "WEBP": "image/webp" } upload_mimetype = mimetype_map.get(chosen_format_upper, 'application/octet-stream') 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') 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,)} NODE_CLASS_MAPPINGS = { "SaveImageToGoogleDrive": SaveImageToGoogleDrive } NODE_DISPLAY_NAME_MAPPINGS = { "SaveImageToGoogleDrive": "Save Image to Google Drive (Creepybits)" }