diff --git a/SaveImageToGoogleDrive.py b/SaveImageToGoogleDrive.py new file mode 100644 index 0000000..852c33f --- /dev/null +++ b/SaveImageToGoogleDrive.py @@ -0,0 +1,246 @@ +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)" +} \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..823f8de --- /dev/null +++ b/__init__.py @@ -0,0 +1,18 @@ +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" + +NODE_CLASS_MAPPINGS = { + **nodes_NODE_CLASS_MAPPINGS, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + **nodes_NODE_DISPLAY_NAME_MAPPINGS, +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..b71bf47 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +google-api-python-client +google-auth-oauthlib +google-auth-httplib2 +Pillow \ No newline at end of file