diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ad8d9c0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.pyc +# Your sign-in token and client ID stay on your machine. +onedrive_msal_token_cache.bin +onedrive_config.json diff --git a/README.md b/README.md index 20a979c..db11e76 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,15 @@ If you don't know how to get the onedrive_msal_token from Microsoft, I have made ___ +### SETUP +Create a file named `onedrive_config.json` in this node's folder (`custom_nodes/ComfyUI-Save_To_OneDrive/`) with your Azure app's client ID: + +```json +{"client_id": "your-application-client-id"} +``` + +Alternatively, set the `ONEDRIVE_CLIENT_ID` environment variable. The first time the node runs, the console shows a link and a code: open the link, enter the code and sign in, and the node continues the upload. After that it stays signed in. + ### INSTRUCTIONS This node and it's code is completely free to use. It should be noted that you will have to get your own **onedrive_msal_token** with Microsoft in order for the node to work preperly. The **onedrive_msal_token** are needed to give the node access to your OneDrive and to create folders as well as save images and videos. @@ -20,6 +29,6 @@ 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**: Images are numbered, so with the prefix "IMG" they are saved as IMG_00001.png, IMG_00002.png and so on (or .jpeg / .webp). If a file with that name already exists in the folder, OneDrive adds a number to the new one instead of overwriting it. **File Format**: Format available is PNG, JPEG and WEBP diff --git a/SaveImageToOneDrive_CreepyBits.py b/SaveImageToOneDrive_CreepyBits.py index 5b4ec80..7a99c96 100644 --- a/SaveImageToOneDrive_CreepyBits.py +++ b/SaveImageToOneDrive_CreepyBits.py @@ -1,28 +1,54 @@ -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 threading +from urllib.parse import quote import msal -CLIENT_ID = YOUR_CLIENt_ID +NODE_DIR = os.path.dirname(os.path.abspath(__file__)) + +# Your Azure app's "Application (client) ID". Set it in onedrive_config.json next to this file +# ({"client_id": "..."}), or with the ONEDRIVE_CLIENT_ID environment variable, or paste it here. +CLIENT_ID = "" +CONFIG_FILE = os.path.join(NODE_DIR, "onedrive_config.json") 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") +TOKEN_CACHE_FILE = os.path.join(NODE_DIR, "onedrive_msal_token_cache.bin") + +UPLOAD_TIMEOUT_SECONDS = 300 + +_msal_app = None +_msal_lock = threading.Lock() -msal_app = msal.PublicClientApplication( - CLIENT_ID, - authority=AUTHORITY -) +def get_client_id(): + client_id = os.environ.get("ONEDRIVE_CLIENT_ID", "").strip() or CLIENT_ID.strip() + if not client_id and os.path.exists(CONFIG_FILE): + try: + with open(CONFIG_FILE, "r", encoding="utf-8") as f: + client_id = str(json.load(f).get("client_id", "")).strip() + except Exception as e: + print(f"[OneDrive Auth] Could not read {CONFIG_FILE}: {e}") + return client_id + + +def get_msal_app(): + """Created on first use rather than at import, so a missing client ID doesn't stop ComfyUI from loading the node. + The token cache is loaded from disk once and then kept in memory.""" + global _msal_app + if _msal_app is None: + client_id = get_client_id() + if not client_id: + return None + _msal_app = msal.PublicClientApplication(client_id, authority=AUTHORITY, token_cache=load_msal_token_cache()) + return _msal_app + def load_msal_token_cache(): cache = msal.SerializableTokenCache() @@ -48,62 +74,68 @@ def save_msal_token_cache(cache): 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 + msal_app = get_msal_app() + if msal_app is None: + message = ("OneDrive client ID is not set. Put {\"client_id\": \"\"} in " + f"{CONFIG_FILE}, or set the ONEDRIVE_CLIENT_ID environment variable.") + print(f"[OneDrive Auth] {message}") + return None, message - accounts = msal_app.get_accounts() - result = None - status_message = None + with _msal_lock: + accounts = msal_app.get_accounts() + result = None + status_message = None - if accounts: - result = msal_app.acquire_token_silent(SCOPES, account=accounts[0]) + if accounts: + result = msal_app.acquire_token_silent(SCOPES, account=accounts[0]) - if not result: - flow = msal_app.initiate_device_flow(scopes=SCOPES) + 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 + 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}") + return None, error_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." + 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"The node waits here until you have signed in, then continues the upload." + ) + print("----------------------------------------------------------------------") + print(device_flow_instructions) + print("----------------------------------------------------------------------") + status_message = device_flow_instructions.replace("\n", " - ") - 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 + 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 + if result and "access_token" in result: + save_msal_token_cache(msal_app.token_cache) + return result["access_token"], "OneDrive authentication successful." + 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}") + return None, error_desc else: + if status_message: + return None, status_message 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: + # Per-prefix counters. OneDrive renames on conflict, so these only need to be unique + # within a session (the temp-folder files they used to be tracked with were wiped on restart anyway). + _counters = {} + _counter_lock = threading.Lock() + def __init__(self): self.output_dir = folder_paths.get_temp_directory() self.type = "temp" @@ -137,33 +169,12 @@ class SaveImageToOneDrive_CreepyBits: 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 get_next_indexed_filename(self, filename_prefix, file_ext): + key = (filename_prefix.lower(), file_ext.lower()) + with self._counter_lock: + next_index = self._counters.get(key, 0) + 1 + self._counters[key] = next_index + return f"{filename_prefix}_{next_index:05d}.{file_ext}" 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() @@ -191,90 +202,78 @@ class SaveImageToOneDrive_CreepyBits: } upload_mimetype = mimetype_map.get(chosen_format_upper, 'application/octet-stream') + # PNG metadata is the same for every image in the batch, so build it once. + 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)) all_uploaded_files = [] all_previews = [] - overall_status_message = "" - temp_actual_image_dir = folder_paths.get_temp_directory() + os.makedirs(self.output_dir, exist_ok=True) - 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) + # One connection for the whole batch instead of a new TLS handshake per image. + with requests.Session() as session: + session.headers.update({"Authorization": f"Bearer {access_token}", "Content-Type": upload_mimetype}) - pil_image_for_save = None - try: - pil_image_for_save = self.tensor_to_pil(image_tensor) + for i, image_tensor in enumerate(images): + indexed_filename_only = self.get_next_indexed_filename(filename_prefix, file_extension) + temp_filepath_local = os.path.join(self.output_dir, indexed_filename_only) - 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) + pil_image_for_save = self.tensor_to_pil(image_tensor) - print(f"[SaveToOneDrive] {upload_status}") + if chosen_format_upper == "JPEG" and pil_image_for_save.mode != "RGB": + pil_image_for_save = pil_image_for_save.convert("RGB") + + if chosen_format_upper == "PNG": + pil_image_for_save.save(temp_filepath_local, pnginfo=metadata, compress_level=4) + else: + pil_image_for_save.save(temp_filepath_local, quality=quality) + + 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}") + all_uploaded_files.append(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 + + # Encode the path so folder names with spaces, '#', '%' etc. upload to the right place. + upload_url = f"https://graph.microsoft.com/v1.0/me/drive/root:/{quote(upload_item_path)}:/content?@microsoft.graph.conflictBehavior=rename" + + try: + response = session.put(upload_url, data=file_bytes, timeout=UPLOAD_TIMEOUT_SECONDS) + response.raise_for_status() + + response_json = response.json() + file_web_url = response_json.get("webUrl", "N/A") + upload_status = f"SUCCESS: Uploaded '{response_json.get('name', 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 Exception: + 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) diff --git a/__init__.py b/__init__.py index 3c24438..531bcbc 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,6 @@ -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" +__version__ = "1.1.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'] diff --git a/pyproject.toml b/pyproject.toml index 23a1e0a..f1b606d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "Save_To_OneDrive" description = "A ComfyUI node that saves images directly to OneDrive using Microsoft's free API service" -version = "1.0.2" +version = "1.1.0" license = {text = "MIT License"} [project.urls] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..ff358b4 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +msal +requests +Pillow +numpy