Add files via upload
This commit is contained in:
Creepybits
@@ -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)"
|
||||
}
|
||||
+18
@@ -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']
|
||||
@@ -0,0 +1,4 @@
|
||||
google-api-python-client
|
||||
google-auth-oauthlib
|
||||
google-auth-httplib2
|
||||
Pillow
|
||||
Reference in New Issue
Block a user