From 6bf9d37204bed9cb2625befdfeb269c06b681bf8 Mon Sep 17 00:00:00 2001 From: Jack Date: Mon, 16 Dec 2024 21:09:09 +0200 Subject: [PATCH] Image Captioning, New API UI with UTM referral, Default 90s timeout + Dynamic timeout, Enhanced Prompt Weighting for txt2img and PhotoMaker, Changed Model Sorting --- __init__.py | 2 +- clientlibs/main.js | 8 +++++- clientlibs/utils.js | 47 +++++++++++++++++++++++++++++++- modules/apiManager.py | 12 +++++++-- modules/bridges/mainRoute.py | 13 +++++++++ modules/imageCaptioning.py | 21 ++++++++++++--- modules/imageInference.py | 20 +++++++------- modules/photoMaker.py | 10 +++++++ modules/utils/runwareUtils.py | 50 ++++++++++++++++++++++++++++------- 9 files changed, 155 insertions(+), 28 deletions(-) diff --git a/__init__.py b/__init__.py index dae5a75..8faa6a0 100644 --- a/__init__.py +++ b/__init__.py @@ -14,7 +14,7 @@ from .modules.apiManager import apiManager from .modules.imageCaptioning import imageCaptioning from .modules.controlNetCombine import controlNetCombine -RUNWARE_COMFYUI_VERSION = "0.9 Beta" +RUNWARE_COMFYUI_VERSION = "0.9.1 Beta" RESET_COLOR = "\033[0m" BLUE_COLOR = "\033[94m" diff --git a/clientlibs/main.js b/clientlibs/main.js index db18f16..3b339f4 100644 --- a/clientlibs/main.js +++ b/clientlibs/main.js @@ -1,12 +1,13 @@ import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; -import { promptEnhanceHandler, syncDimensionsNodeHandler, searchNodeHandler, APIKeyHandler, handleCustomErrors } from "./utils.js"; +import { promptEnhanceHandler, syncDimensionsNodeHandler, searchNodeHandler, APIKeyHandler, captionNodeHandler, handleCustomErrors } from "./utils.js"; import { RUNWARE_NODE_TYPES, RUNWARE_NODE_PROPS, SEARCH_TERMS } from "./types.js"; app.registerExtension({ name: "runware.ai", async setup() { api.addEventListener('runwareError', handleCustomErrors); + api.addEventListener('runwareImageCaption', captionNodeHandler); }, async nodeCreated(node) { @@ -21,6 +22,11 @@ app.registerExtension({ node.bgcolor = crNodeProps.bgColor; if(nodeClass === RUNWARE_NODE_TYPES.APIMANAGER) { APIKeyHandler(node); + } else if(nodeClass === RUNWARE_NODE_TYPES.IMAGECAPTION) { + const captionInput = node.widgets[1].inputEl; + captionInput.style.outline = "none"; + captionInput.readOnly = true; + return; } if(crNodeProps.colorModeOnly === true) return; diff --git a/clientlibs/utils.js b/clientlibs/utils.js index 5bffdaf..f165e94 100644 --- a/clientlibs/utils.js +++ b/clientlibs/utils.js @@ -5,6 +5,7 @@ import { DEFAULT_DIMENSIONS_LIST, DEFAULT_MODELS_ARCH_LIST, DEFAULT_CONTROLNET_C } from "./types.js"; let openDialog = false; +let lastTimeout = false; async function enhancePrompt(userPrompt) { const resp = await api.fetchApi('/promptEnhance', { @@ -30,6 +31,18 @@ async function setAPIKey(apiKey) { return setAPIKeyResults; } +async function setTimeout(maxTimeout) { + const resp = await api.fetchApi('/setMaxTimeout', { + method: "POST", + headers: {'Content-Type': 'application/json'}, + body: JSON.stringify({ + maxTimeout: maxTimeout + }) + }); + const setTimeoutResults = await resp.json(); + return setTimeoutResults; +} + async function modelSearch(modelQuery = "", modelArch = "all", modelType = "base", modelCat = "checkpoint", condtioning = "") { const resp = await api.fetchApi('/modelSearch', { method: "POST", @@ -46,6 +59,20 @@ async function modelSearch(modelQuery = "", modelArch = "all", modelType = "base return modelSearchResults; } +function captionNodeHandler(msgEvent) { + const captionData = msgEvent.detail; + const captionText = captionData.captionText; + const captionNodeID = parseInt(captionData.nodeID); + if(captionData.success) { + const captionNode = app.graph.getNodeById(captionNodeID); + if(captionNode !== null || captionNode !== undefined) { + const captionInput = captionNode.widgets[1].inputEl; + captionInput.value = captionText; + } + } + return false; +} + function notifyUser(message, type="info", title = "Runware", life = 4.5) { app.extensionManager.toast.add({ severity: type, // 'info', 'success', 'warn', 'error' \\ @@ -99,7 +126,24 @@ async function APIKeyHandler(apiManagerNode) { } } }); - break; + } else if(widget.name === "Max Timeout") { + if(lastTimeout) widget.value = lastTimeout; + appendWidgetCB(widget, async function(...args) { + const maxTimeout = args[0]; + if(maxTimeout < 5 || maxTimeout > 99) { + notifyUser("Invalid Timeout Value, Timeout must be between 5 and 99 seconds!", "error", "Runware API Manager"); + return; + } else { + const setTimeoutResults = await setTimeout(maxTimeout); + if(setTimeoutResults.success) { + remNode(apiManagerNode); + notifyUser("Timeout Set Successfully!", "success", "Runware API Manager"); + lastTimeout = maxTimeout; + } else { + notifyUser(setTimeoutResults.error, "error", "Runware API Manager"); + } + } + }); } } } @@ -319,6 +363,7 @@ export { promptEnhanceHandler, syncDimensionsNodeHandler, searchNodeHandler, + captionNodeHandler, handleCustomErrors, APIKeyHandler }; \ No newline at end of file diff --git a/modules/apiManager.py b/modules/apiManager.py index a266871..42b9942 100644 --- a/modules/apiManager.py +++ b/modules/apiManager.py @@ -1,14 +1,22 @@ +from .utils import runwareUtils as rwUtils + class apiManager: @classmethod def INPUT_TYPES(cls): return { "required": { "API Key": ("STRING", { - "tooltip": "Set 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", + "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": "Define Or Change Runware API Timeout In Seconds.\n\n(e.g; 90 Seconds).", + "min": 5, + "max": 99, + "default": int(rwUtils.getTimeout()), }), }, } - DESCRIPTION = "API Managers is a Runware Utility That helps you define or change your API Key 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 or Session Timeout 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 33ee558..d72d003 100644 --- a/modules/bridges/mainRoute.py +++ b/modules/bridges/mainRoute.py @@ -22,6 +22,18 @@ async def setAPIKey(reqPayload): return web.json_response({'success': False, 'error': str(e)}) return web.json_response({'success': True}) +@routes.post('/setMaxTimeout') +async def setMaxTimeout(reqPayload): + reqData = await reqPayload.json() + maxTimeout = reqData.get('maxTimeout', 90) + if(maxTimeout < 5 or maxTimeout > 99): + return web.json_response({'success': False, 'error': 'Invalid Timeout Value!'}) + try: + rwUtils.setTimeout(maxTimeout) + 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() @@ -55,6 +67,7 @@ async def modelSearch(reqPayload): "taskUUID": rwUtils.genRandUUID(), "category": modelCategory, "limit": 15, + "sort": "-downloadCount", }] if(modelCategory != "controlnet" and modelCategory != "lora"): diff --git a/modules/imageCaptioning.py b/modules/imageCaptioning.py index 12e74ac..931c292 100644 --- a/modules/imageCaptioning.py +++ b/modules/imageCaptioning.py @@ -5,30 +5,42 @@ class imageCaptioning: def INPUT_TYPES(cls): return { "required": { + "Always Recaption": ("BOOLEAN", { + "default": False, + "tooltip": "Enable this option to always recaption the image each time you run the workflow.", + "label_on": "Enabled", + "label_off": "Disabled", + }), "Image": ("IMAGE", { - "tooltip": "Specifies the input image to be processed." + "tooltip": "Specify The Image to be Captioned." }), }, "optional": { "Image Description": ("STRING", { "multiline": True, - "placeholder": "Don't Write Anything Here, The Image Description Will Be Generated Automatically.", + "placeholder": "You Don't Have to write Anything here.\nThe Image Description Will Be Generated Automatically.", }), }, "hidden": { "node_id": "UNIQUE_ID" } } DESCRIPTION = "Image to text, also known as image captioning, allows you to obtain descriptive text prompts based on uploaded or previously generated images. This process is instrumental in generating textual descriptions that can be used to create additional images or provide detailed insights into visual content." + FUNCTION = "imageCaptioning" RETURN_TYPES = ("STRING",) RETURN_NAMES = ("IMAGE PROMPT",) CATEGORY = "Runware" OUTPUT_NODE = True + @classmethod + def IS_CHANGED(s, **kwargs): + alwaysRecaption = kwargs.get("Always Recaption") + if(alwaysRecaption): + return float("NAN") + return True + def imageCaptioning(self, **kwargs): - print(kwargs) image = kwargs.get("Image") - kwargs.set("Image Description", "Don't Write Anything Here, The Image Description Will Be Generated Automatically.") genConfig = [ { "taskType": "imageCaption", @@ -39,4 +51,5 @@ class imageCaptioning: genResult = rwUtils.inferenecRequest(genConfig) genText = genResult["data"][0]["text"] + rwUtils.sendImageCaption(genText, kwargs.get("node_id")) return (genText, ) \ No newline at end of file diff --git a/modules/imageInference.py b/modules/imageInference.py index f1ca5ec..f700e6c 100644 --- a/modules/imageInference.py +++ b/modules/imageInference.py @@ -18,11 +18,9 @@ class txt2img: "placeholder": "Negative Prompt: a text instruction to guide the model on generating the image. It is usually a sentence or a paragraph that provides negative guidance for the task. This parameter helps to avoid certain undesired results.", "tooltip": "Negative Prompt: a text instruction to guide the model on generating the image." }), - "Prompt Weighting": ("BOOLEAN", { - "default": False, - "tooltip": "Prompt weighting allows you to adjust how strongly different parts of your prompt influence the generated image.", - "label_on": "Enabled", - "label_off": "Disabled", + "Prompt Weighting": (["Disabled", "sdEmbeds", "Compel"], { + "default": "Disabled", + "tooltip": "Prompt weighting allows you to adjust how strongly different parts of your prompt influence the generated image.\n\nChoose between \"compel\" notation with advanced weighting operations or \"sdEmbeds\" for simple emphasis adjustments.\n\nCompel Example: \"small+ dog, pixar style\"\n\nsdEmbeds Example: \"(small:2.5) dog, pixar style\"", }), "dimensions": ([ "Square (512x512)", "Square HD (1024x1024)", "Portrait 3:4 (768x1024)", @@ -140,7 +138,7 @@ class txt2img: runwareModel = kwargs.get("Model") positivePrompt = kwargs.get("positivePrompt") negativePrompt = kwargs.get("negativePrompt", None) - promptWeighting = kwargs.get("Prompt Weighting", False) + promptWeighting = kwargs.get("Prompt Weighting", "Disabled") runwareControlNet = kwargs.get("ControlNet", None) runwareLora = kwargs.get("Lora", None) runwareRefiner = kwargs.get("Refiner", None) @@ -185,8 +183,11 @@ class txt2img: if(negativePrompt is not None and negativePrompt != ""): genConfig[0]["negativePrompt"] = negativePrompt - if(promptWeighting): - genConfig[0]["promptWeighting"] = "sdEmbeds" + if(promptWeighting != "Disabled"): + if(promptWeighting == "sdEmbeds"): + genConfig[0]["promptWeighting"] = "sdEmbeds" + else: + genConfig[0]["promptWeighting"] = "compel" if(runwareLora is not None): if(isinstance(runwareLora, list)): genConfig[0]["lora"] = runwareLora @@ -208,4 +209,5 @@ class txt2img: genResult = rwUtils.inferenecRequest(genConfig) images = rwUtils.convertImageB64List(genResult) - return images \ No newline at end of file + return images + \ No newline at end of file diff --git a/modules/photoMaker.py b/modules/photoMaker.py index e6cd54a..0f6a6ca 100644 --- a/modules/photoMaker.py +++ b/modules/photoMaker.py @@ -21,6 +21,10 @@ class photoMaker: "placeholder": "Negative Prompt: a text instruction to guide the model on generating the image. It is usually a sentence or a paragraph that provides negative guidance for the task. This parameter helps to avoid certain undesired results.", "tooltip": "Negative Prompt: a text instruction to guide the model on generating the image." }), + "Prompt Weighting": (["Disabled", "Compel"], { + "default": "Disabled", + "tooltip": "Prompt weighting allows you to adjust how strongly different parts of your prompt influence the generated image.", + }), "style": ([ "No style", "Cinematic", "Disney Character", "Digital Art", "Photographic", "Fantasy art", "Neonpunk", "Enhance", "Comic book", "Lowpoly", "Line art" @@ -129,6 +133,7 @@ class photoMaker: image4 = kwargs.get("Image 4", None) positivePrompt = kwargs.get("positivePrompt") negativePrompt = kwargs.get("negativePrompt", None) + promptWeighting = kwargs.get("Prompt Weighting", "Disabled") style = kwargs.get("style", "No style") dimensions = kwargs.get("dimensions", "Square HD (1024x1024)") width = kwargs.get("width", 1024) @@ -173,6 +178,11 @@ class photoMaker: if(negativePrompt is not None and negativePrompt != ""): genConfig[0]["negativePrompt"] = negativePrompt + if(promptWeighting != "Disabled"): + if(promptWeighting == "sdEmbeds"): + genConfig[0]["promptWeighting"] = "sdEmbeds" + else: + genConfig[0]["promptWeighting"] = "compel" genResult = rwUtils.inferenecRequest(genConfig) images = rwUtils.convertImageB64List(genResult) diff --git a/modules/utils/runwareUtils.py b/modules/utils/runwareUtils.py index 14ab78e..e40d610 100644 --- a/modules/utils/runwareUtils.py +++ b/modules/utils/runwareUtils.py @@ -15,7 +15,6 @@ import io load_dotenv() -SESSION_TIMEOUT = 60 RUNWARE_API_BASE_URL = "https://api.runware.ai/v1" session = requests.Session() adapter = HTTPAdapter(pool_connections=10, pool_maxsize=10) @@ -28,10 +27,19 @@ def getAPIKey(): return apiKey return False +def getTimeout(): + timeout = os.getenv("RUNWARE_TIMEOUT") + if timeout and timeout.isdigit(): + return int(timeout) + else: + timeout = 90 + os.environ["RUNWARE_TIMEOUT"] = str(timeout) + return timeout + +SESSION_TIMEOUT = getTimeout() RUNWARE_API_KEY = getAPIKey() -def setAPIKey(apiKey: str): - global RUNWARE_API_KEY +def setEnvKey(keyName, keyValue): comfyNodeRoot = Path(__file__).parent.parent.parent envFilePath = comfyNodeRoot / '.env' if not envFilePath.exists(): @@ -41,19 +49,33 @@ def setAPIKey(apiKey: str): key_exists = False new_lines = [] for line in lines: - if line.startswith('RUNWARE_API_KEY='): + if line.startswith(f'{keyName}='): key_exists = True - new_lines.append(f'RUNWARE_API_KEY={apiKey}\n') + new_lines.append(f'{keyName}={keyValue}\n') else: new_lines.append(line) if not key_exists: - new_lines.append(f'RUNWARE_API_KEY={apiKey}\n') + new_lines.append(f'{keyName}={keyValue}\n') with open(envFilePath, 'w') as f: f.writelines(new_lines) - RUNWARE_API_KEY = apiKey - os.environ["RUNWARE_API_KEY"] = apiKey return True +def setAPIKey(apiKey: str): + global RUNWARE_API_KEY + envSetRes = setEnvKey("RUNWARE_API_KEY", apiKey) + if envSetRes: + RUNWARE_API_KEY = apiKey + os.environ["RUNWARE_API_KEY"] = apiKey + return True + +def setTimeout(timeout: int): + envSetRes = setEnvKey("RUNWARE_TIMEOUT", str(timeout)) + if envSetRes: + global SESSION_TIMEOUT + SESSION_TIMEOUT = timeout + os.environ["RUNWARE_TIMEOUT"] = str(timeout) + return True + def genRandSeed(minSeed = 1000, maxSeed = 9223372036854776000): return random.randint(minSeed, maxSeed) @@ -73,7 +95,7 @@ def checkAPIKey(apiKey): }] try: - genResult = session.post(RUNWARE_API_BASE_URL, headers=headers, json=genConfig, timeout=SESSION_TIMEOUT, allow_redirects=False, stream=True) + genResult = session.post(RUNWARE_API_BASE_URL, headers=headers, json=genConfig, timeout=10, allow_redirects=False, stream=True) genResult = genResult.json() if("errors" in genResult): return str(genResult["errors"][0]["message"]) @@ -82,9 +104,17 @@ def checkAPIKey(apiKey): except Exception as e: return False +def sendImageCaption(captionText, nodeID): + PromptServer.instance.send_sync("runwareImageCaption", { + "success": True, + "captionText": captionText, + "nodeID": nodeID, +}) + def inferenecRequest(genConfig): - global RUNWARE_API_KEY + global RUNWARE_API_KEY, RUNWARE_API_BASE_URL, SESSION_TIMEOUT RUNWARE_API_KEY = os.getenv("RUNWARE_API_KEY") + SESSION_TIMEOUT = int(os.getenv("RUNWARE_TIMEOUT")) headers = { "Authorization": f"Bearer {RUNWARE_API_KEY}", "Content-Type": "application/json",