From 8a4418c9a92fee6768a3cf9057cfe603ae6664d5 Mon Sep 17 00:00:00 2001 From: Creepybits Date: Tue, 15 Jul 2025 14:40:12 +0200 Subject: [PATCH] Add files via upload --- SaveImageToOneDrive_CreepyBits.py | 293 ++++++++++++++++++++++++++++++ __init__.py | 18 ++ 2 files changed, 311 insertions(+) create mode 100644 SaveImageToOneDrive_CreepyBits.py create mode 100644 __init__.py diff --git a/SaveImageToOneDrive_CreepyBits.py b/SaveImageToOneDrive_CreepyBits.py new file mode 100644 index 0000000..5b4ec80 --- /dev/null +++ b/SaveImageToOneDrive_CreepyBits.py @@ -0,0 +1,293 @@ +import torch +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo +import io +import requests +import os +import folder_paths +import json +import re +import time + +import msal + +CLIENT_ID = YOUR_CLIENt_ID +AUTHORITY = "https://login.microsoftonline.com/consumers" +SCOPES = ["https://graph.microsoft.com/Files.ReadWrite", "https://graph.microsoft.com/User.Read"] + +TOKEN_CACHE_FILE = os.path.join(os.path.dirname(__file__), "onedrive_msal_token_cache.bin") + + +msal_app = msal.PublicClientApplication( + CLIENT_ID, + authority=AUTHORITY +) + +def load_msal_token_cache(): + cache = msal.SerializableTokenCache() + if os.path.exists(TOKEN_CACHE_FILE): + try: + with open(TOKEN_CACHE_FILE, "r") as f: + cache.deserialize(f.read()) + except Exception as e: + print(f"[OneDrive Auth] Error loading MSAL token cache: {e}. Starting with an empty cache.") + try: + os.remove(TOKEN_CACHE_FILE) + except OSError as oe: + print(f"[OneDrive Auth] Error removing corrupted token cache file: {oe}") + return cache + +def save_msal_token_cache(cache): + if cache.has_state_changed: + try: + with open(TOKEN_CACHE_FILE, "w") as f: + f.write(cache.serialize()) + cache.has_state_changed = False + except Exception as e: + print(f"[OneDrive Auth] Error saving MSAL token cache: {e}") + +def get_onedrive_access_token_msal(): + current_token_cache = load_msal_token_cache() + msal_app.token_cache = current_token_cache + + accounts = msal_app.get_accounts() + result = None + status_message = None + + if accounts: + result = msal_app.acquire_token_silent(SCOPES, account=accounts[0]) + + if not result: + flow = msal_app.initiate_device_flow(scopes=SCOPES) + + if "user_code" not in flow: + error_message = f"Failed to create device flow. Error: {flow.get('error_description', str(flow))}" + print(f"[OneDrive Auth] {error_message}") + status_message = error_message + return None, status_message + + device_flow_instructions = ( + f"ONEDRIVE AUTH NEEDED (Device Flow):\n" + f"1. Open browser to: {flow['verification_uri']}\n" + f"2. Enter code: {flow['user_code']}\n" + f"3. Sign in & grant permissions.\n" + f"Node will then attempt to get token. Re-queue prompt after authorizing." + ) + print("----------------------------------------------------------------------") + print(device_flow_instructions) + print("----------------------------------------------------------------------") + status_message = device_flow_instructions.replace("\n", " - ") + " - Check console for details and re-queue." + + try: + result = msal_app.acquire_token_by_device_flow(flow) + except Exception as e_device_flow: + print(f"[OneDrive Auth] Error during acquire_token_by_device_flow: {e_device_flow}") + status_message = f"Error during device flow: {e_device_flow}" + result = None + + if result and "access_token" in result: + save_msal_token_cache(msal_app.token_cache) + status_message = "OneDrive authentication successful." + return result["access_token"], status_message + elif result and "error_description" in result: + error_desc = f"Failed to acquire token. Error: {result.get('error')}, Desc: {result.get('error_description')}" + print(f"[OneDrive Auth] {error_desc}") + status_message = error_desc + return None, status_message + else: + if status_message: + return None, status_message + else: + unknown_error_msg = "[OneDrive Auth] Could not acquire token. Unknown error or device flow not completed." + print(unknown_error_msg) + return None, unknown_error_msg + +class SaveImageToOneDrive_CreepyBits: + def __init__(self): + self.output_dir = folder_paths.get_temp_directory() + self.type = "temp" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", ), + "filename_prefix": ("STRING", {"default": "ComfyUI"}), + "onedrive_folder_path": ("STRING", {"default": "ComfyUI_Uploads"}), + "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 = ("STRING",) + RETURN_NAMES = ("status_message",) + OUTPUT_NODE = True + FUNCTION = "execute" + CATEGORY = "Creepybits/Cloud" + + def tensor_to_pil(self, tensor_image): + if tensor_image.ndim == 4: + tensor_image = tensor_image[0] + img_np = tensor_image.cpu().numpy() * 255.0 + pil_img = Image.fromarray(np.clip(img_np, 0, 255).astype(np.uint8)) + return pil_img + + def get_next_indexed_filename(self, output_dir, filename_prefix, file_ext): + index_tracker_dir = os.path.join(output_dir, "onedrive_node_index_tracker") + os.makedirs(index_tracker_dir, exist_ok=True) + + pattern = re.compile(rf"^{re.escape(filename_prefix)}_(\d+)\.{re.escape(file_ext)}$", re.IGNORECASE) + + max_index = 0 + try: + for f in os.listdir(index_tracker_dir): + match = pattern.match(f) + if match: + current_index = int(match.group(1)) + if current_index > max_index: + max_index = current_index + except FileNotFoundError: + pass + + next_index = max_index + 1 + indexed_filename_only = f"{filename_prefix}_{next_index:05d}.{file_ext}" + + try: + with open(os.path.join(index_tracker_dir, indexed_filename_only), 'wb') as f: + f.write(b'') + except Exception as e: + print(f"[SaveToOneDrive] Warning: Could not create dummy index file {indexed_filename_only} in {index_tracker_dir}: {e}") + + return indexed_filename_only + + def execute(self, images, onedrive_folder_path, filename_prefix, file_format, quality, prompt=None, extra_pnginfo=None): + access_token, auth_status_message = get_onedrive_access_token_msal() + + if not access_token: + print(f"[SaveToOneDrive] Authentication failed or requires user action. Message: {auth_status_message}") + display_msg = auth_status_message.replace("\n", " - ") if auth_status_message else \ + "AUTH_ERROR: OneDrive authentication failed. Check console for details and re-queue prompt." + return {"ui": {"string": [display_msg]}, "result": (display_msg,)} + + cleaned_folder_path = onedrive_folder_path.strip().strip('/') + + chosen_format_upper = file_format.upper() + file_extension = "png" + + if chosen_format_upper == "JPEG": + file_extension = "jpeg" + elif chosen_format_upper == "WEBP": + file_extension = "webp" + + mimetype_map = { + "PNG": "image/png", + "JPEG": "image/jpeg", + "WEBP": "image/webp" + } + upload_mimetype = mimetype_map.get(chosen_format_upper, 'application/octet-stream') + + + all_uploaded_files = [] + all_previews = [] + overall_status_message = "" + + temp_actual_image_dir = folder_paths.get_temp_directory() + + for i, image_tensor in enumerate(images): + indexed_filename_only = self.get_next_indexed_filename( + self.output_dir, filename_prefix, file_extension + ) + temp_filepath_local = os.path.join(temp_actual_image_dir, indexed_filename_only) + + pil_image_for_save = None + try: + pil_image_for_save = self.tensor_to_pil(image_tensor) + + if chosen_format_upper == "JPEG": + if pil_image_for_save.mode == 'RGBA': + pil_image_for_save = pil_image_for_save.convert("RGB") + elif chosen_format_upper == "WEBP": + if pil_image_for_save.mode == 'P': + pil_image_for_save = pil_image_for_save.convert("RGBA") + elif pil_image_for_save.mode == 'RGBA' and not pil_image_for_save.getchannel('A').getbbox(): + pil_image_for_save = pil_image_for_save.convert("RGB") + + 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)) + + os.makedirs(temp_actual_image_dir, exist_ok=True) + if chosen_format_upper == "PNG": + pil_image_for_save.save(temp_filepath_local, pnginfo=metadata, compress_level=4) + elif chosen_format_upper in ["JPEG", "WEBP"]: + pil_image_for_save.save(temp_filepath_local, quality=quality) + else: + pil_image_for_save.save(temp_filepath_local) + + with open(temp_filepath_local, 'rb') as f: + file_bytes = f.read() + + except Exception as e: + error_msg = f"ERROR processing image {i+1} for local save: {str(e)}" + print(f"[SaveToOneDrive] {error_msg}") + overall_status_message += f"Image {i+1} Failed: {error_msg}. " + continue + + upload_item_path = f"{cleaned_folder_path}/{indexed_filename_only}" if cleaned_folder_path else indexed_filename_only + + upload_url = f"https://graph.microsoft.com/v1.0/me/drive/root:/{upload_item_path}:/content?@microsoft.graph.conflictBehavior=rename" + + headers = { + "Authorization": f"Bearer {access_token}", + "Content-Type": upload_mimetype + } + + try: + response = requests.put(upload_url, headers=headers, data=file_bytes) + response.raise_for_status() + + response_json = response.json() + file_web_url = response_json.get("webUrl", "N/A") + upload_status = f"SUCCESS: Uploaded '{indexed_filename_only}'. URL: {file_web_url}" + all_uploaded_files.append(upload_status) + + all_previews.append({"filename": indexed_filename_only, "subfolder": "", "type": "temp"}) + + except requests.exceptions.HTTPError as http_err: + error_details = str(http_err.response.text if hasattr(http_err.response, 'text') and http_err.response.text else http_err) + try: + error_details = http_err.response.json().get("error", {}).get("message", error_details) + except: pass + upload_status = f"ERROR: HTTP {http_err.response.status_code} uploading '{indexed_filename_only}': {error_details}" + all_uploaded_files.append(upload_status) + except Exception as e: + upload_status = f"ERROR: Unexpected error during upload of '{indexed_filename_only}': {str(e)}" + all_uploaded_files.append(upload_status) + + print(f"[SaveToOneDrive] {upload_status}") + + if all_uploaded_files: + overall_status_message = " | ".join(all_uploaded_files) + else: + overall_status_message = "No images were uploaded due to errors." + + ui_return_data = {} + if all_previews: + ui_return_data["images"] = all_previews + if overall_status_message: + ui_return_data["string"] = [overall_status_message] + + return {"ui": ui_return_data, "result": (overall_status_message,)} + +NODE_CLASS_MAPPINGS = { "SaveImageToOneDrive_CreepyBits": SaveImageToOneDrive_CreepyBits } +NODE_DISPLAY_NAME_MAPPINGS = { "SaveImageToOneDrive_CreepyBits": "Save Image to OneDrive (Creepybits)" } diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3c24438 --- /dev/null +++ b/__init__.py @@ -0,0 +1,18 @@ +import os +import sys +from .SaveImageToOneDrive_CreepyBits 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']