diff --git a/.gitignore b/.gitignore index e1777d9..a49e3ec 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,7 @@ __pycache__/ *.py[cod] *$py.class - +imagesCache.json .env .venv diff --git a/clientlibs/utils.js b/clientlibs/utils.js index 500961c..bb0ba12 100644 --- a/clientlibs/utils.js +++ b/clientlibs/utils.js @@ -6,11 +6,14 @@ import { DEFAULT_DIMENSIONS_LIST, DEFAULT_MODELS_ARCH_LIST, DEFAULT_CONTROLNET_C const TIMEOUT_RANGE = { min: 5, default: 90, max: 99 }; const OUTPUT_QUALITY_RANGE = { min: 20, default: 95, max: 99 }; +const CACHE_SIZE_RANGE = { min: 30, default: 100, max: 4096 }; let openDialog = false; let lastTimeout = false; let lastOutputFormat = false; let lastOutputQuality = false; +let lastEnableImagesCaching = null; +let lastMinImageCacheSize = false; async function queryLocalAPI(endpoint, data) { try { @@ -31,6 +34,8 @@ const runwareLocalAPI = { setTimeout: (maxTimeout) => queryLocalAPI('setMaxTimeout', { maxTimeout }), setOutputFormat: (outputFormat) => queryLocalAPI('setOutputFormat', { outputFormat }), setOutputQuality: (outputQuality) => queryLocalAPI('setOutputQuality', { outputQuality }), + setEnableImagesCaching: (enableCaching) => queryLocalAPI('setEnableImagesCaching', { enableCaching }), + setMinImageCacheSize: (minCacheSize) => queryLocalAPI('setMinImageCacheSize', { minCacheSize }), modelSearch: (modelQuery = "", modelArch = "all", modelType = "base", modelCat = "checkpoint", condtioning = "") => queryLocalAPI('modelSearch', { modelQuery, modelArch, modelType, modelCat, condtioning }) }; @@ -149,6 +154,37 @@ async function APIKeyHandler(apiManagerNode) { notifyUser(setQualityResults.error, "error", "Runware API Manager"); } }); + } else if(widget.name === "Enable Images Caching") { + if(lastEnableImagesCaching !== null) widget.value = lastEnableImagesCaching; + appendWidgetCB(widget, async function(...args) { + const enableCaching = args[0]; + const setCachingResults = await runwareLocalAPI.setEnableImagesCaching(enableCaching); + if(setCachingResults.success) { + notifyUser("Image Caching Setting Updated Successfully!", "success", "Runware API Manager"); + lastEnableImagesCaching = enableCaching; + } else { + widget.value = lastEnableImagesCaching; + notifyUser(setCachingResults.error, "error", "Runware API Manager"); + } + }); + } else if(widget.name === "Min Image Cache Size") { + if(lastMinImageCacheSize !== false) widget.value = lastMinImageCacheSize; + appendWidgetCB(widget, async function(...args) { + const minCacheSize = parseFloat(args[0]); + if(isNaN(minCacheSize) || minCacheSize < CACHE_SIZE_RANGE.min || minCacheSize > CACHE_SIZE_RANGE.max) { + notifyUser(`Invalid cache size! Must be between ${CACHE_SIZE_RANGE.min} KB and ${CACHE_SIZE_RANGE.max} KB.`, "error", "Runware API Manager"); + widget.value = lastMinImageCacheSize || CACHE_SIZE_RANGE.default; + return; + } + const setCacheSizeResults = await runwareLocalAPI.setMinImageCacheSize(minCacheSize); + if(setCacheSizeResults.success) { + notifyUser("Minimum Image Cache Size Updated Successfully!", "success", "Runware API Manager"); + lastMinImageCacheSize = minCacheSize; + } else { + widget.value = lastMinImageCacheSize || CACHE_SIZE_RANGE.default; + notifyUser(setCacheSizeResults.error, "error", "Runware API Manager"); + } + }); } } } diff --git a/modules/apiManager.py b/modules/apiManager.py index 7038956..924cf3e 100644 --- a/modules/apiManager.py +++ b/modules/apiManager.py @@ -5,28 +5,59 @@ class apiManager: def INPUT_TYPES(cls): return { "required": { - "API Key": ("STRING", { - "tooltip": "Define Or Change Your Runware API Key To Start Using Runware Image Inference Services.\n\n(e.g; BcVr1JVQM8FGzrRwb3GsbCq5ww1QwabV)\n\nIf You Don't Have One, You Can Get It From: https://my.runware.ai/keys", - }), - "Max Timeout": ("INT", { - "tooltip": "Change Runware API Request Timeout In Seconds.\n\n(e.g; 90 Seconds).", - "min": 5, - "max": 99, - "default": int(rwUtils.getTimeout()), - }), - "Image Output Quality": ("INT", { - "tooltip": "Sets the compression quality of the output image. Higher values preserve more quality but increase file size, lower values reduce file size but decrease quality.", - "min": 20, - "max": 99, - "default": int(rwUtils.getOutputQuality()), - }), - "Image Output Format": (["WEBP", "PNG", "JPEG"], { - "tooltip": "Change the Default Image Output Format.", - "default": rwUtils.getOutputFormat(), - }), + "API Key": ( + "STRING", + { + "tooltip": "Define Or Change Your Runware API Key To Start Using Runware Image Inference Services.\n\n(e.g; BcVr1JVQM8FGzrRwb3GsbCq5ww1QwabV)\n\nIf You Don't Have One, You Can Get It From: https://my.runware.ai/keys", + }, + ), + "Max Timeout": ( + "INT", + { + "tooltip": "Change Runware API Request Timeout In Seconds.\n\n(e.g; 90 Seconds).", + "min": 5, + "max": 99, + "default": int(rwUtils.getTimeout()), + }, + ), + "Image Output Quality": ( + "INT", + { + "tooltip": "Sets the compression quality of the output image. Higher values preserve more quality but increase file size, lower values reduce file size but decrease quality.", + "min": 20, + "max": 99, + "default": int(rwUtils.getOutputQuality()), + }, + ), + "Image Output Format": ( + ["WEBP", "PNG", "JPEG"], + { + "tooltip": "Change the Default Image Output Format.", + "default": rwUtils.getOutputFormat(), + }, + ), + "Enable Images Caching": ( + "BOOLEAN", + { + "label_on": "Enabled", + "label_off": "Disabled", + "tooltip": "Enable or disable image caching functionality.\n\nWhen enabled, images will be cached to improve performance and reduce redundant processing.", + "default": rwUtils.getEnableImagesCaching(), + }, + ), + "Min Image Cache Size": ( + "INT", + { + "tooltip": "Set the minimum size (in KB) for images to be cached.", + "min": 30, + "max": 4096, + "step": 1, + "default": int(rwUtils.getMinImageCacheSize()), + }, + ), }, } - DESCRIPTION = "API Managers is a Runware Utility That helps you define or change your API Key, Session Timeout, or Image Output Format & Quality From The ComfyUI Interface Without having to adjust the env file locally." + DESCRIPTION = "API Managers is a Runware Utility That helps you define or change your API Key, Session Timeout, Image Output Format & Quality, and Image Caching settings From The ComfyUI Interface Without having to adjust the env file locally." CATEGORY = "Runware" RETURN_TYPES = () \ No newline at end of file diff --git a/modules/bridges/mainRoute.py b/modules/bridges/mainRoute.py index 375af8a..17055d0 100644 --- a/modules/bridges/mainRoute.py +++ b/modules/bridges/mainRoute.py @@ -58,6 +58,30 @@ async def setOutputQuality(reqPayload): return web.json_response({'success': False, 'error': str(e)}) return web.json_response({'success': True}) +@routes.post('/setEnableImagesCaching') +async def setEnableImagesCaching(reqPayload): + reqData = await reqPayload.json() + enableCaching = reqData.get('enableCaching', True) + if not isinstance(enableCaching, bool): + return web.json_response({'success': False, 'error': 'Invalid value for enable caching!'}) + try: + rwUtils.setEnableImagesCaching(enableCaching) + except Exception as e: + return web.json_response({'success': False, 'error': str(e)}) + return web.json_response({'success': True}) + +@routes.post('/setMinImageCacheSize') +async def setMinImageCacheSize(reqPayload): + reqData = await reqPayload.json() + minCacheSize = reqData.get('minCacheSize', 100) + if not isinstance(minCacheSize, int) or minCacheSize < 30 or minCacheSize > 4096: + return web.json_response({'success': False, 'error': 'Invalid minimum cache size!'}) + try: + rwUtils.setMinImageCacheSize(minCacheSize) + except Exception as e: + return web.json_response({'success': False, 'error': str(e)}) + return web.json_response({'success': True}) + @routes.post('/promptEnhance') async def promptEnhance(reqPayload): reqData = await reqPayload.json() diff --git a/modules/utils/runwareUtils.py b/modules/utils/runwareUtils.py index e526b3b..fa93696 100644 --- a/modules/utils/runwareUtils.py +++ b/modules/utils/runwareUtils.py @@ -1,26 +1,39 @@ from comfy.model_management import InterruptProcessingException from requests.adapters import HTTPAdapter +from datetime import datetime, timedelta from server import PromptServer from dotenv import load_dotenv from pathlib import Path from PIL import Image import numpy as np import requests +import hashlib +import asyncio import random import base64 import torch import uuid import time +import json import os import io +import threading load_dotenv() +BASEFOLDER = Path(__file__).parent.parent.parent +IMAGE_CACHE_FILE = BASEFOLDER / "imagesCache.json" + +if not IMAGE_CACHE_FILE.exists(): + with open(IMAGE_CACHE_FILE, "w") as f: + print("[Runware] Initializing images cache...") + json.dump({}, f) + RUNWARE_REMBG_OUTPUT_FORMATS = { - "outputFormat": (["WEBP", "PNG"], { - "default": "WEBP", - "tooltip": "Choose the output image format." - }) + "outputFormat": ( + ["WEBP", "PNG"], + {"default": "WEBP", "tooltip": "Choose the output image format."}, + ) } RUNWARE_API_BASE_URL = "https://api.runware.ai/v1" @@ -77,29 +90,49 @@ def getOutputFormat(): os.environ["RUNWARE_OUTPUT_FORMAT"] = output_format return output_format +def getEnableImagesCaching(): + enable_images_caching = os.getenv("RUNWARE_ENABLE_IMAGES_CACHING") + if enable_images_caching and enable_images_caching.lower() in ["true", "false"]: + return enable_images_caching.lower() == "true" + else: + enable_images_caching = True + os.environ["RUNWARE_ENABLE_IMAGES_CACHING"] = str(enable_images_caching) + return enable_images_caching + +def getMinImageCacheSize(): + min_image_cache_size = os.getenv("RUNWARE_MIN_IMAGE_CACHE_SIZE") + if min_image_cache_size and min_image_cache_size.isdigit(): + return int(min_image_cache_size) + else: + min_image_cache_size = 100 + os.environ["RUNWARE_MIN_IMAGE_CACHE_SIZE"] = str(min_image_cache_size) + return min_image_cache_size + SESSION_TIMEOUT = getTimeout() RUNWARE_API_KEY = getAPIKey() OUTPUT_FORMAT = getOutputFormat() OUTPUT_QUALITY = getOutputQuality() +ENABLE_IMAGES_CACHING = getEnableImagesCaching() +MIN_IMAGE_CACHE_SIZE = getMinImageCacheSize() def setEnvKey(keyName, keyValue): comfyNodeRoot = Path(__file__).parent.parent.parent - envFilePath = comfyNodeRoot / '.env' + envFilePath = comfyNodeRoot / ".env" if not envFilePath.exists(): envFilePath.touch() - with open(envFilePath, 'r') as f: + with open(envFilePath, "r") as f: lines = f.readlines() key_exists = False new_lines = [] for line in lines: - if line.startswith(f'{keyName}='): + if line.startswith(f"{keyName}="): key_exists = True - new_lines.append(f'{keyName}={keyValue}\n') + new_lines.append(f"{keyName}={keyValue}\n") else: new_lines.append(line) if not key_exists: - new_lines.append(f'{keyName}={keyValue}\n') - with open(envFilePath, 'w') as f: + new_lines.append(f"{keyName}={keyValue}\n") + with open(envFilePath, "w") as f: f.writelines(new_lines) return True @@ -135,7 +168,23 @@ def setOutputQuality(quality: int): os.environ["RUNWARE_OUTPUT_QUALITY"] = str(quality) return True -def genRandSeed(minSeed = 1000, maxSeed = 9223372036854776000): +def setEnableImagesCaching(enabled: bool): + envSetRes = setEnvKey("RUNWARE_ENABLE_IMAGES_CACHING", str(enabled)) + if envSetRes: + global ENABLE_IMAGES_CACHING + ENABLE_IMAGES_CACHING = enabled + os.environ["RUNWARE_ENABLE_IMAGES_CACHING"] = str(enabled) + return True + +def setMinImageCacheSize(size: int): + envSetRes = setEnvKey("RUNWARE_MIN_IMAGE_CACHE_SIZE", str(size)) + if envSetRes: + global MIN_IMAGE_CACHE_SIZE + MIN_IMAGE_CACHE_SIZE = size + os.environ["RUNWARE_MIN_IMAGE_CACHE_SIZE"] = str(size) + return True + +def genRandSeed(minSeed=1000, maxSeed=9223372036854776000): return random.randint(minSeed, maxSeed) def genRandUUID(): @@ -148,18 +197,28 @@ def checkAPIKey(apiKey): "Accept-Encoding": "gzip, deflate, br, zstd", } - genConfig = [{ + genConfig = [ + { "taskType": "authentication", "apiKey": apiKey, - }] + } + ] try: + def recaller(): - return session.post(RUNWARE_API_BASE_URL, headers=headers, json=genConfig, timeout=10, allow_redirects=False, stream=True) - + return session.post( + RUNWARE_API_BASE_URL, + headers=headers, + json=genConfig, + timeout=10, + allow_redirects=False, + stream=True, + ) + genResult = generalRequestWrapper(recaller) genResult = genResult.json() - if("errors" in genResult): + if "errors" in genResult: return str(genResult["errors"][0]["message"]) else: return True @@ -167,11 +226,14 @@ def checkAPIKey(apiKey): return False def sendImageCaption(captionText, nodeID): - PromptServer.instance.send_sync("runwareImageCaption", { - "success": True, - "captionText": captionText, - "nodeID": nodeID, -}) + PromptServer.instance.send_sync( + "runwareImageCaption", + { + "success": True, + "captionText": captionText, + "nodeID": nodeID, + }, + ) def inferenecRequest(genConfig): global RUNWARE_API_KEY, RUNWARE_API_BASE_URL, SESSION_TIMEOUT @@ -184,39 +246,137 @@ def inferenecRequest(genConfig): } try: + def recaller(): - return session.post(RUNWARE_API_BASE_URL, headers=headers, json=genConfig, timeout=SESSION_TIMEOUT, allow_redirects=False, stream=True) - + return session.post( + RUNWARE_API_BASE_URL, + headers=headers, + json=genConfig, + timeout=SESSION_TIMEOUT, + allow_redirects=False, + stream=True, + ) + genResult = generalRequestWrapper(recaller) genResult = genResult.json() - if("errors" in genResult): + if "errors" in genResult: raise Exception(genResult["errors"][0]["message"]) else: return genResult except requests.exceptions.Timeout: - raise Exception(f"Error: Request Timed Out After {SESSION_TIMEOUT} Seconds - Please Try Again!") + raise Exception( + f"Error: Request Timed Out After {SESSION_TIMEOUT} Seconds - Please Try Again!" + ) except requests.exceptions.RequestException: raise Exception("Error: Runware Request Failed!") except Exception as e: if "invalid api key" in str(e).lower(): - PromptServer.instance.send_sync("runwareError", { - "success": False, - "errorMessage": str(e), - "errorCode": 401, - }) + PromptServer.instance.send_sync( + "runwareError", + { + "success": False, + "errorMessage": str(e), + "errorCode": 401, + }, + ) raise InterruptProcessingException() else: raise Exception(f"Error: {e}") return False + +async def uploadImage(imageDataUri): + uploadTaskConfig = [ + {"taskType": "imageUpload", "taskUUID": genRandUUID(), "image": imageDataUri} + ] + + try: + uploadResult = inferenecRequest(uploadTaskConfig) + if ( + uploadResult + and "data" in uploadResult + and "imageUUID" in uploadResult["data"][0] + ): + imageUUID = uploadResult["data"][0]["imageUUID"] + return imageUUID + except Exception as e: + return False + + return False + +async def uploadAndCacheImage(imgSig: str, imgDataUri: str): + try: + imageUUID = await uploadImage(imgDataUri) + if imageUUID: + await imageStoreSet(imgSig, imageUUID) + except Exception as e: + return False + +async def imageStoreSet(imgHash: str, imgUUID: str) -> bool: + try: + with open(IMAGE_CACHE_FILE, "r") as f: + cache = json.load(f) + expires = (datetime.now() + timedelta(days=30)).isoformat() + cache[imgHash] = {"uuid": imgUUID, "expires": expires} + with open(IMAGE_CACHE_FILE, "w") as f: + json.dump(cache, f, indent=2) + return True + except Exception as e: + return False + +def imageStoreGet(imgHash: str) -> str | bool: + try: + with open(IMAGE_CACHE_FILE, "r") as f: + cache = json.load(f) + if imgHash not in cache: + return False + entry = cache[imgHash] + expires = datetime.fromisoformat(entry["expires"]) + if datetime.now() > expires: + del cache[imgHash] + with open(IMAGE_CACHE_FILE, "w") as f: + json.dump(cache, f, indent=2) + return False + return entry["uuid"] + except Exception as e: + return False + def convertTensor2IMG(tensorImage): + global ENABLE_IMAGES_CACHING, MIN_IMAGE_CACHE_SIZE + ENABLE_IMAGES_CACHING = getEnableImagesCaching() + MIN_IMAGE_CACHE_SIZE = int(os.getenv("RUNWARE_MIN_IMAGE_CACHE_SIZE")) + imageNP = (tensorImage.squeeze().numpy() * 255).astype(np.uint8) + imgSig = hashlib.sha256(imageNP.tobytes()).hexdigest() + + if ENABLE_IMAGES_CACHING: + imgUUID = imageStoreGet(imgSig) + if imgUUID: + return imgUUID + image = Image.fromarray(imageNP) with io.BytesIO() as buffer: - image.save(buffer, format="webp") - imgb64 = base64.b64encode(buffer.getvalue()).decode('utf-8') - imgDataURI = f'data:image/webp;base64,{imgb64}' - return imgDataURI + image.save( + buffer, format="webp", quality=100, subsampling=0, method=6, exact=True + ) + imageRawData = buffer.getvalue() + imgBytes = len(imageRawData) + imgSize = int(imgBytes / 1024) + imgb64 = base64.b64encode(imageRawData).decode("utf-8") + imgDataUri = f"data:image/webp;base64,{imgb64}" + + if ENABLE_IMAGES_CACHING and imgSize >= MIN_IMAGE_CACHE_SIZE: + try: + loop = asyncio.get_running_loop() + loop.call_soon(lambda: asyncio.ensure_future(uploadAndCacheImage(imgSig, imgDataUri))) + except RuntimeError: + def run_coro(): + new_loop = asyncio.new_event_loop() + asyncio.set_event_loop(new_loop) + new_loop.run_until_complete(uploadAndCacheImage(imgSig, imgDataUri)) + new_loop.close() + threading.Thread(target=run_coro).start() + return imgDataUri def convertIMG2Tensor(b64img): imgbytes = base64.b64decode(b64img) @@ -228,8 +388,13 @@ def convertIMG2Tensor(b64img): def convertImageB64List(imageDataObject): images = () for result in imageDataObject["data"]: - generatedImage = result.get("imageBase64Data", result.get("maskImageBase64Data", result.get("guideImageBase64Data", False))) + generatedImage = result.get( + "imageBase64Data", + result.get( + "maskImageBase64Data", result.get("guideImageBase64Data", False) + ), + ) generatedImage = convertIMG2Tensor(generatedImage) images += (generatedImage,) images = torch.stack(images, dim=0) - return (images,) \ No newline at end of file + return (images,)