Added New Image Caching System

This commit is contained in:
Jack
2025-02-22 00:15:04 +02:00
parent b82ed10f4b
commit 27f5d69b90
5 changed files with 313 additions and 57 deletions
+1 -1
View File
@@ -1,7 +1,7 @@
__pycache__/
*.py[cod]
*$py.class
imagesCache.json
.env
.venv
+36
View File
@@ -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");
}
});
}
}
}
+51 -20
View File
@@ -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 = ()
+24
View File
@@ -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()
+201 -36
View File
@@ -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,)
return (images,)