commit 36df15595f0fe69a6c235cdb607dbfacde86173a Author: Beamhi Date: Sat Aug 24 08:10:27 2024 -0700 first commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..933c844 --- /dev/null +++ b/.gitignore @@ -0,0 +1,164 @@ +mass-process/output +mass-process/image +mass-process/*.png + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000..251a8a5 --- /dev/null +++ b/README.md @@ -0,0 +1,98 @@ +# ComfyUI_NetDist Plus +Run ComfyUI workflows on multiple local GPUs/networked machines with options to edit the json values within comfyui. +# Credits +[City96](https://github.com/city96); without the base netdist repo, I wouldn't have attempted this. +[EventStationAI](https://www.eventstation.ai/); for some GPU support. +[Claude](https://console.anthropic.com/dashboard); what do I do next? Can you debug this error? +[Ogkai](https://civitai.com/user/ogkai_1111); Thanks for encouraging me to start pushing stuffs I make or modify. + + + +[NetDist_2xspeed.webm](https://github.com/city96/ComfyUI_NetDist/assets/125218114/b7ec2fcf-1e51-4b05-ad62-355da2a1bf6d) + +## Install instructions: +There is currently a single external requirement, which is the `requests` library. +``` +pip install requests +``` + +To install, simply clone into the custom nodes folder. +``` +git clone https://github.com/city96/ComfyUI_NetDist ComfyUI/custom_nodes/ComfyUI_NetDist +``` + +## Usage + +### Local Remote control +You will need at least two different ComfyUI instances. You can use two local GPUs by setting different `--port [port]` and `--cuda-device [number]` launch arguments. You'll most likely want `--port 8288 --cuda-device 1` + +#### Simple dual-GPU + +This is the simplest setup for people who have 2 GPUs or two separate PCs. It only requires two nodes to work. + +You can set the local/remote batch size, as well as when the node should trigger (set it to 'always' if it isn't getting executed - i.e. you changed a sampler setting but not the seed.) + +If you're running your second instance on a different PC, add `--listen` to your launch arguments and set the correct remote IP (open a terminal window and check with `ipconfig` on windows or `ip a` on linux). + +The `FetchRemote` ('Fetch from remote') node takes an image input. This should be your final image than you want to get back from your second instance (make sure not to route it back into itself). This node will wait for the second image to be generated (there's currently no preview/progress bar). + +Workflow JSON: [NetDistSimple.json](https://github.com/city96/ComfyUI_NetDist/files/13825326/NetDistSimple.json) + +![NetDistSimple](https://github.com/city96/ComfyUI_NetDist/assets/125218114/dce5a155-2ffa-4979-b184-03de168beecb) + +#### Simple multi-machine + +You can kind of scale the example above by connecting more of the simple queue nodes together, but the seed is a bit jank and you can get duplicate images if you try and reuse it. I guess just set the seed to randomized on both. + +![NetDistMulti](https://github.com/city96/ComfyUI_NetDist/assets/125218114/2a0358aa-ab8e-47e2-82a2-7a27a17d0130) + +#### Advanced + +This is mostly meant for more "advanced" setups with more than two GPUs. It allows easier per-batch overrides as well as setting a default batch size. + +It also allows using a workflow JSON as an input. To allow any workflow to run, the final image can be set to "any" instead of the default "final_image" (which would require the `FetchRemote` node to be in the workflow). + +I have nodes to save/load the workflows, but ideally there would be some nodes to also edit them - search and replace seed, etc. PRs welcome ;P + +Workflow JSON: [NetDistAdvancedV2.json](https://github.com/city96/ComfyUI_NetDist/files/13843005/NetDistAdvancedV2.json) + +![NetDistAdvanced](https://github.com/city96/ComfyUI_NetDist/assets/125218114/851c1ee6-edcf-4489-bab1-92ab9c5ef15e) + +(This needs a fake image input to trigger, you can just give it a blank image). + +![NetDistSaved](https://github.com/city96/ComfyUI_NetDist/assets/125218114/a39b5117-af1b-4f2c-a94e-5a330acc8ea4) + +### Remote images +The `LoadImageUrl` ('Load Image (URL)') Node acts just like the normal 'Load Image' node. + +The `SaveImageUrl` ('Save Image (URL)') Node sends a POST request to the target URL with a json containing the images. +- The filenames are the keys. +- The values are the base64 encoded PNG images (optionally with the `data:image/png;base64` prefix). +- The filenames are **not** guaranteed to be unique across batches since they aren't saved locally. You should handle this server-side. +- No data is written to disk on the server. + +### Remote latents + +This node pack has a set of nodes which should (in theory) allow you to pass latents between the nodes seamlessly. A node to save the input latent as a `.npy` file is provided. This node also returns the filename of the saved latent, which can then be loaded by the other instance. + +To load a latent from the other instance, you can plug the filename into this URL: + +``` +# change the filename with a string replacement node. +http://127.0.0.1:8188/view?filename=ComfyUI_00001_.latent&type=output` +# To load them from the input folder instead, change type to 'input' +http://127.0.0.1:8188/view?filename=TestLatent.npy&type=input +``` + +The `LoadLatentNumpy` node can also load the default safetensor latents, the npy ones (simple numpy file containing just the latent in the standard torch format) as well as the sd_scripts npz cache files. + +![LatentSave](https://github.com/city96/ComfyUI_NetDist/assets/125218114/cd68d8dc-bd96-4018-82c9-400337fc5f80) + +### Things you probably shouldn't do: +- Queue a workflow on the same remote worker multiple times from the same client. +- ~~Expect this to work smoothly.~~ + +## Roadmap +- Fix some edge cases, like linux controlling windows (`os.sep` mismatch). +- Better workflow editing for static workflows. +- Handle multiple separate image output nodes. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..a706110 --- /dev/null +++ b/__init__.py @@ -0,0 +1,25 @@ +# only import if running as a custom node +try: + import comfy.utils +except ImportError: + pass +else: + NODE_CLASS_MAPPINGS = {} + + from .nodes.simple import NODE_CLASS_MAPPINGS as NetNodes + NODE_CLASS_MAPPINGS.update(NetNodes) + + from .nodes.advanced import NODE_CLASS_MAPPINGS as AdvNodes + NODE_CLASS_MAPPINGS.update(AdvNodes) + + from .nodes.images import NODE_CLASS_MAPPINGS as ImgNodes + NODE_CLASS_MAPPINGS.update(ImgNodes) + + from .nodes.latents import NODE_CLASS_MAPPINGS as LatNodes + NODE_CLASS_MAPPINGS.update(LatNodes) + + from .nodes.workflows import NODE_CLASS_MAPPINGS as WrkNodes + NODE_CLASS_MAPPINGS.update(WrkNodes) + + NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()} + __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/core/dispatch.py b/core/dispatch.py new file mode 100644 index 0000000..01d677c --- /dev/null +++ b/core/dispatch.py @@ -0,0 +1,160 @@ +import os +import time +import json +import torch +import random +import requests +import numpy as np +from PIL import Image +from copy import deepcopy + +from .utils import clean_url, get_client_id + +def clear_remote_queue(remote_url): + r = requests.get(f"{remote_url}/queue", timeout=4) + r.raise_for_status() + queue = r.json() + + to_cancel = [] + client_id = get_client_id() + for k in queue.get("queue_pending", []): + if k[3].get("client_id") == client_id: + to_cancel.append(k[1]) # job UUID + r = requests.post( + f"{remote_url}/queue", + json = {"delete" : to_cancel}, + timeout = 4, + ) + r.raise_for_status() + + for k in queue.get("queue_running", []): + if k[3].get("client_id") == client_id: + r = requests.post( + f"{remote_url}/interrupt", + json = {}, + timeout = 4, + ) + r.raise_for_status() + break + +def get_remote_os(remote_url): + url = f"{remote_url}/system_stats" + r = requests.get(url, timeout=4) + r.raise_for_status() + data = r.json() + return data["system"]["os"] + +def get_output_nodes(remote_url): + # I'm 90% sure this could just use the + # list from the host but better safe than sorry + url = f"{remote_url}/object_info" + r = requests.get(url, timeout=4) + r.raise_for_status() + data = r.json() + out = [k for k, v in data.items() if v.get("output_node")] + return out + + +def dispatch_to_remote(remote_url, prompt, job_id=f"{get_client_id()}-unknown", remote_params=[], outputs="final_image"): + ### PROMPT LOGIC ### + prompt = deepcopy(prompt) + to_del = [] + + def recursive_node_deletion(start_node): + target_nodes = [start_node] + if start_node not in to_del: + to_del.append(start_node) + while len(target_nodes) > 0: + new_targets = [] + for target in target_nodes: + for node in prompt.keys(): + inputs = prompt[node].get("inputs") + if not inputs: + continue + for i in inputs.values(): + if type(i) == list: + if len(i) > 0 and i[0] in to_del: + if node not in to_del: + to_del.append(node) + new_targets.append(node) + target_nodes += new_targets + target_nodes.remove(target) + + # find current node and disable all others + output_src = None + for i in prompt.keys(): + if prompt[i]["class_type"].startswith("RemoteQueue"): + if clean_url(prompt[i]["inputs"]["remote_url"]) == remote_url: + prompt[i]["inputs"]["enabled"] = "remote" + output_src = i + # Apply remote parameters + for param, value, nodeid in remote_params: + if nodeid: + for node_key, node_data in prompt.items(): + if node_key == nodeid: + if param in node_data.get("inputs", {}): + prompt[node_key]["inputs"][param] = value + break + else: + if param in prompt[i]["inputs"]: + prompt[i]["inputs"][param] = value + else: + # If the parameter doesn't exist in the node's inputs, + # we need to find where to apply it in the prompt + for node_key, node_data in prompt.items(): + if param in node_data.get("inputs", {}): + prompt[node_key]["inputs"][param] = value + break + else: + prompt[i]["inputs"]["enabled"] = "false" + + banned = [] if outputs == "any" else ["PreviewImage", "SaveImage"] # get_output_nodes(remote_url) + output = None + for i in prompt.keys(): + # only leave current fetch but replace with PreviewImage + if prompt[i]["class_type"] == "FetchRemote": + if prompt[i]["inputs"]["remote_info"][0] == output_src: + output = { + "inputs": {"images": prompt[i]["inputs"]["final_image"]}, + "class_type": 'PreviewImage', + "final_output": True, # might allow multiple outputs with an ID? + } + recursive_node_deletion(i) + # do not save output on remote + if prompt[i]["class_type"] in banned: + recursive_node_deletion(i) + if output: + prompt[str(max([int(x) for x in prompt.keys()])+1)] = output + for i in to_del: del prompt[i] + + ### OS LOGIC ### + sep_remote = "\\" if get_remote_os(remote_url) == "nt" else "/" + sep_local = "\\" if os.name == "nt" else "/" + sem_input_map = { # class type : input to replace + "CheckpointLoaderSimple" : "ckpt_name", + "CheckpointLoader" : "ckpt_name", + "LoraLoader" : "lora_name", + "VAELoader" : "vae_name", + } + if sep_remote != sep_local: + for i in prompt.keys(): + if prompt[i]["class_type"] in sem_input_map.keys(): + key = sem_input_map[prompt[i]["class_type"]] + prompt[i]["inputs"][key] = prompt[i]["inputs"][key].replace(sep_local, sep_remote) + + ### SEND REQUEST ### + data = { + "prompt": prompt, + "client_id": get_client_id(), + "extra_data": { + "job_id": job_id, + } + } + ar = requests.post( + f"{remote_url}/prompt", + data = json.dumps(data), + headers = {"Content-Type": "application/json"}, + timeout = 4, + ) + ar.raise_for_status() + return diff --git a/core/fetch.py b/core/fetch.py new file mode 100644 index 0000000..41f9cb1 --- /dev/null +++ b/core/fetch.py @@ -0,0 +1,97 @@ +import time +import json +import torch +import requests +import numpy as np +from PIL import Image + +POLLING = 0.5 + +def get_job_output(inputs, outputs): + output_id = list(outputs.keys())[-1] # fallback to last + for i,d in inputs.items(): + if d.get("final_output") and i in outputs.keys(): + output_id = i + break + return outputs[output_id].get("images", []) + +def wait_for_job(remote_url, job_id): + fail = 0 + while fail <= 3: + r = requests.get(f"{remote_url}/history", timeout=4) + try: + r.raise_for_status() + except Exception as e: + print("NetDist caught error while fetching output image:\n", e) + fail += 1 + continue + data = r.json() + if not data: + time.sleep(POLLING) + continue + for i,d in data.items(): + if d["prompt"][3].get("job_id") == job_id: + # this needs to be less jank + if len(d["outputs"].keys()) > 0: + return get_job_output(d["prompt"][2], d["outputs"]) + else: + return [] + # todo: check if it's actually in the queue to avoid waiting forever + time.sleep(POLLING) + raise OSError("Failed to fetch image from remote client!") + +def fetch_from_remote(remote_url, job_id): + def img_to_torch(img): + image = img.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image)[None,] + image_tensor.metadata = img.info # Store metadata in tensor attribute + return image_tensor + + if not remote_url or not job_id: + return None + + images = [] + for i in wait_for_job(remote_url, job_id): + img_url = f"{remote_url}/view?filename={i['filename']}&subfolder={i['subfolder']}&type={i['type']}" + + ir = requests.get(img_url, stream=True, timeout=16) + ir.raise_for_status() + img = Image.open(ir.raw) + images.append(img_to_torch(img)) + + if len(images) == 0: + return None + + out = images[0] + for i in images[1:]: + out = torch.cat((out, i)) + return out + +#with extras returns both the output and the metadata from the images generated remotely +def fetch_from_remote_with_extras(remote_url, job_id): + def img_to_torch(img): + image = img.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image)[None,] + return image_tensor + + if not remote_url or not job_id: + return None + + images = [] + for i in wait_for_job(remote_url, job_id): + img_url = f"{remote_url}/view?filename={i['filename']}&subfolder={i['subfolder']}&type={i['type']}" + + ir = requests.get(img_url, stream=True, timeout=16) + ir.raise_for_status() + img = Image.open(ir.raw) + images.append(img_to_torch(img)) + + if len(images) == 0: + return None + + out = images[0] + for i in images[1:]: + out = torch.cat((out, i)) + return out, img.info \ No newline at end of file diff --git a/core/utils.py b/core/utils.py new file mode 100644 index 0000000..07e9737 --- /dev/null +++ b/core/utils.py @@ -0,0 +1,23 @@ +import time +import random + +# set global ID once for entire session +try: GID +except NameError: + GID = ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) + print(f"NetDist: Set session ID to '{GID}'") + +def get_client_id(): + global GID + return(f"netdist-{GID}") + +def get_new_job_id(): + job_id = f"{get_client_id()}-{int(time.time()*1000)}" + time.sleep(0.1) # prevent ID mismatch, no matter how unlikely + return job_id + +def clean_url(raw, multi=False): + raw = raw.strip() + raw = raw.replace(' ', ',').replace('\n', ',').replace('\t', ',') + urls = [x.rstrip('/') for x in raw.split(',') if x.strip()] + return urls if multi else urls[0] diff --git a/fonts/Roboto-Regular.ttf b/fonts/Roboto-Regular.ttf new file mode 100644 index 0000000..3033308 Binary files /dev/null and b/fonts/Roboto-Regular.ttf differ diff --git a/mass-process/job.example.yaml b/mass-process/job.example.yaml new file mode 100644 index 0000000..68f8962 --- /dev/null +++ b/mass-process/job.example.yaml @@ -0,0 +1,28 @@ +workflow: "job.example.png" # Your actual workflow to distribute. +job_start: 1 # Start of job_num, will be incremented +job_end: 100 # until it reaches job_end. + +workers: # List of workers the server will connect to. + "RTX3080@LOC": # Client nickname. Can be anything. + url: "http://127.0.0.1:8188/" # ComfyUI URL (will be used for API). + system: "windows" # 'windows' or 'posix'. + "P40_0@NET": + url: "http://192.168.4.6:8188/" + system: "posix" + "P40_1@NET": + url: "http://192.168.4.6:8288/" + system: "posix" + +# Replace specific strings in the workflow inputs. +# Source string should be the one present in your saved workflow. +# Destination string will be the one the clients see. +replacement: + - + src: "http://127.0.0.1:8080/image/openpose/0000.png" + dst: "http://127.0.0.1:8080/image/openpose/{job_num:04}.png" + - + src: "http://192.168.4.4:8080/0000.png" + dst: "http://192.168.4.4:8080/{job_num:04}.png" + - + src: "http://example.lan/upload/output/0000.png" + dst: "http://example.lan/upload/output/{job_num:04}.png" diff --git a/mass-process/readme.md b/mass-process/readme.md new file mode 100644 index 0000000..ccfd59b --- /dev/null +++ b/mass-process/readme.md @@ -0,0 +1,9 @@ +This is the script I used when I had to mass-process some controlnet inputs for an animation. It was requested that I publish this in [issue#2](https://github.com/city96/ComfyUI_NetDist/issues/2#issuecomment-1696342450) + +Note that you'll have to make all images accessible to the clients, otherwise they fail. I was using the load from URL nodes for this. For testing, I simply used the built-in python web server with `python -m http.server 8080`, but you can use your own web server to host them. + +I don't currently have seed randomization or proper output handling. For the later, just disable all the save/preview image nodes other than the one you'll be using as the final output. + +Here's a sample workflow: + +![job example](https://github.com/city96/ComfyUI_NetDist/assets/125218114/138ec97b-61a6-4631-a280-06b5c0e3c43d) diff --git a/mass-process/server.py b/mass-process/server.py new file mode 100644 index 0000000..33e38f0 --- /dev/null +++ b/mass-process/server.py @@ -0,0 +1,159 @@ +import os +import time +import yaml +import json +import requests +import argparse +from PIL import Image +from tqdm import tqdm +from queue import Queue +from copy import deepcopy +from threading import Thread + +class JobShard: + def __init__(self, workflow, job_num): + self.workflow = workflow # raw workflow + self.job_num = job_num # numerical ID of job + self.prompt = None # created when assigned to worker + self.job_id = None # ^ + + def format_workflow(self, rep, system, job_num): + w = deepcopy(self.workflow) + for i in w.keys(): + # Fix path mismatch + ct = w[i]["class_type"] + pr = ("\\","/") if system == "posix" else ("/","\\") + if ct == "LoraLoader": + w[i]["inputs"]["lora_name"] = w[i]["inputs"]["lora_name"].replace(*pr) + elif ct == "VAELoader": + w[i]["inputs"]["vae_name"] = w[i]["inputs"]["vae_name"].replace(*pr) + elif ct in ["CheckpointLoader","CheckpointLoaderSimple"]: + w[i]["inputs"]["ckpt_name"] = w[i]["inputs"]["ckpt_name"].replace(*pr) + # replace strings + for k in w[i].get("inputs",{}).keys(): + src = w[i]["inputs"][k] + dst = [x["dst"] for x in rep if x["src"] == src] + if dst: + w[i]["inputs"][k] = dst[0].format(job_num=job_num) + self.prompt = w + + def assign(self, worker): + self.format_workflow(worker.conf["replacement"], worker.system, self.job_num) + self.job_id = f"{worker.name}-{self.job_num}@{int(time.time())}" + +class Worker: + def __init__(self, name, system, url, conf, jobs, prog): + self.name = name + self.url = url.rstrip("/") if url.endswith("/") else url + self.system = system.lower().strip() + self.conf = conf # global config + self.jobs = jobs # queue of all jobs + self.prog = prog # progress bar + self.job = None + + def is_busy(self): + busy = True if self.job else False + return busy + + def run(self): + while not self.jobs.empty(): + self.job = self.jobs.get() + self.job.assign(self) + self.start_job() + self.fetch_job() + self.job = None + self.jobs.task_done() + self.prog.update() + + def start_job(self): + url = f"{self.url}/prompt" + data = { + "prompt": self.job.prompt, + "client_id": "netdist-mass", + "extra_data": { + "job_id": self.job.job_id, + } + } + r = requests.post(url, json=data) + r.raise_for_status() + + def wait_for_job(self): + url = self.url + "/history" + image_data = None + while not image_data: + r = requests.get(url) + r.raise_for_status() + data = r.json() + if not data: + time.sleep(0.5) + continue + for i,d in data.items(): + if d["prompt"][3].get("job_id") == self.job.job_id: + image_data = d["outputs"][list(d["outputs"].keys())[-1]].get("images") + break + time.sleep(0.5) + return image_data + + def fetch_job(self): + images = [] + for i in self.wait_for_job(): + img_url = f"{self.url}/view?filename={i['filename']}&subfolder={i['subfolder']}&type={i['type']}" + ir = requests.get(img_url, stream=True) + ir.raise_for_status() + images.append(Image.open(ir.raw)) + + if len(images) == 0: + print(f"{self.name}@{self.url} job failed") + elif len(images) == 1: + images[0].save(f"output/{self.job.job_num}.png") + else: + for i in range(len(images)): + images[i].save(f"output/{self.job.job_num}.{i}.png") + +def get_workflow(path): + if path.endswith(".png"): + img = Image.open(path) + data = json.loads( + img.text.get("prompt") + ) + else: + print("invalid input workflow") + exit(1) + return data + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('--conf', required=True, help="Config file describing job.") + args = parser.parse_args() + + with open(args.conf) as f: + conf = yaml.safe_load(f.read()) + + if not os.path.isdir("output"): + os.mkdir("output") + + # creqte queue with jobs + jobs = Queue() + wf = get_workflow(conf["workflow"]) + for job_num in range(conf["job_start"],conf["job_end"]): + jobs.put( + JobShard(wf, job_num) + ) + prog = tqdm(total=jobs.qsize()) + + # initialize workers + workers = [] + for name, k in conf["workers"].items(): + workers.append(Worker( + name=name, + system=k["system"], + url=k["url"], + jobs=jobs, + prog=prog, + conf=conf) + ) + + + # execute all + [Thread(target=w.run, daemon=True).start() for w in workers] + jobs.join() diff --git a/nodes/advanced.py b/nodes/advanced.py new file mode 100644 index 0000000..d970512 --- /dev/null +++ b/nodes/advanced.py @@ -0,0 +1,371 @@ +from ..core.utils import clean_url, get_client_id, get_new_job_id +from ..core.dispatch import dispatch_to_remote, clear_remote_queue + +import copy + +class RemoteApplyValues: + """Apply values to remote nodes""" + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "nodeid": ("STRING", {"default": ""}), + "param": ("STRING", {"default": ""}), + "value": ("STRING", {"default": ""}), + "type": (["STRING", "INT", "FLOAT", "BOOL"], {"default": "STRING"}), + } + } + + RETURN_TYPES = ("REMOTEAPPLY",) + RETURN_NAMES = ("remote_apply",) + FUNCTION = "apply_values" + CATEGORY = "remote/advanced" + TITLE = "Apply Values to Remote Nodes" + + def apply_values(self, nodeid, param, value, type): + return ((nodeid, param, value, type),) + +class RemoteApplyValuesMulti: + """Apply values to remote nodes""" + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "nodeid": ("STRING", {"default": ""}), + "param": ("STRING", {"default": ""}), + "value": ("STRING", {"default": ""}), + "type": ("STRING", {"default": "STRING"}), + } + } + + RETURN_TYPES = ("REMOTEAPPLY",) + RETURN_NAMES = ("remote_apply",) + FUNCTION = "apply_values" + CATEGORY = "remote/advanced" + TITLE = "Apply Values to Remote Nodes" + + def apply_values(self, nodeid, param, value, type): + nodeids = [id.strip() for id in nodeid.split(",")] + params = [p.strip() for p in param.split(",")] + values = [v.strip() for v in value.split(",")] + types = type.split(",") # Split the types string into a list + + # If only one type is provided, apply it to all values + if len(types) == 1: + types = types * len(nodeids) + else: + types = [t.strip() for t in types] # Strip whitespace from each type + + if len(nodeids) != len(params) or len(nodeids) != len(values) or len(nodeids) != len(types): + raise ValueError("Number of nodeids, params, values, and types must be equal") + + remote_values = tuple(zip(nodeids, params, values, types)) + return (remote_values,) + +class RemoteChainStartNux: + """Merge required attributes into one [REMCHAIN]""" + def __init__(self): + pass + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "workflow": ("JSON",), + "trigger": (["on_change", "always"],), + "batch": ("INT", {"default": 1, "min": 1, "max": 8}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + }, + "optional": { + "remoteapply1": ("REMOTEAPPLY",), + "remoteapply2": ("REMOTEAPPLY",), + "remoteapply3": ("REMOTEAPPLY",), + "remoteapply4": ("REMOTEAPPLY",), + "remoteapply5": ("REMOTEAPPLY",), + "remoteapply6": ("REMOTEAPPLY",), + "remoteapply7": ("REMOTEAPPLY",), + "remoteapply8": ("REMOTEAPPLY",), + "remoteapply9": ("REMOTEAPPLY",), + "remoteapply10": ("REMOTEAPPLY",), + } + } + + RETURN_TYPES = ("REMCHAIN",) + RETURN_NAMES = ("remote_chain",) + FUNCTION = "chain_start" + CATEGORY = "remote/advanced" + TITLE = "Queue on remote (start of chain)" + + def chain_start(self, workflow, trigger, batch, seed, + remoteapply1=None, remoteapply2=None, remoteapply3=None, remoteapply4=None, + remoteapply5=None, remoteapply6=None, remoteapply7=None, remoteapply8=None, + remoteapply9=None, remoteapply10=None): + + # Make a deep copy of the workflow to avoid retaining state between calls + workflow = copy.deepcopy(workflow) + + remote_params = {} + for remoteapply in [remoteapply1, remoteapply2, remoteapply3, remoteapply4, + remoteapply5, remoteapply6, remoteapply7, remoteapply8, + remoteapply9, remoteapply10]: + print("**"+str(remoteapply6)+str(type(remoteapply6))+"**") + if remoteapply: + if isinstance(remoteapply[0], tuple): # Check if it's a tuple of tuples (multi) + for nodeid, param, value, value_type in remoteapply: + if param and value: + + remote_params[(nodeid, param)] = self.parse_value(value, value_type) + else: # Single tuple + nodeid, param, value, value_type = remoteapply + if param and value: + remote_params[(nodeid, param)] = self.parse_value(value, value_type) + + remote_params[("", "seed")] = self.parse_value(seed, "INT") # Add seed to remote_params + + # Apply remote parameters to the prompt + for (nodeid, param), value in remote_params.items(): + truncated_value = str(value)[:30] + '...' if len(str(value)) > 30 else str(value) # Truncate value for display + print(f"Applying param: {param}, value: {truncated_value}, nodeid: {nodeid}") # Debug statement + if nodeid: + if nodeid in workflow: + if param in workflow[nodeid].get("inputs", {}): + workflow[nodeid]["inputs"][param] = value + print(f"Updated node {nodeid} param {param} with value {truncated_value}") # Debug statement + else: + print(f"Param {param} not found in node {nodeid} inputs") # Debug statement + else: + print(f"Node {nodeid} not found in workflow") # Debug statement + else: + for node_key, node_data in workflow.items(): + if param in node_data.get("inputs", {}): + workflow[node_key]["inputs"][param] = value + print(f"Updated node {node_key} param {param} with value {truncated_value}") # Debug statement + break + else: + print(f"Param {param} not found in node {node_key} inputs") # Debug statement + + remote_chain = { + "seed": seed, + "batch": batch, + "prompt": workflow, + "seed_offset": batch, + "job_id": get_new_job_id(), + } + return(remote_chain,) + + @classmethod + def IS_CHANGED(self, workflow, trigger, batch, seed, prompt): + uuid = f"W:{workflow},B:{batch},S:{seed}" + return uuid if trigger == "on_change" else str(time.time()) + + def parse_value(self, value, value_type): + """Parse the value string to the specified type.""" + if value_type == "INT": + return int(value) + elif value_type == "FLOAT": + return float(value) + elif value_type == "BOOL": + return value.lower() == "true" + else: # STRING or any other type + return value.replace("\\", "\\\\") + +class RemoteChainStart: + """Merge required attributes into one [REMCHAIN]""" + def __init__(self): + pass + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "workflow": ("JSON",), + "trigger": (["on_change", "always"],), + "batch": ("INT", {"default": 1, "min": 1, "max": 8}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + }, + "optional": { + "remote_nodeid1": ("STRING", {"default": ""}), + "remote_param1": ("STRING", {"default": ""}), + "remote_value1": ("STRING", {"default": ""}), + "remote_type1": (["STRING", "INT", "FLOAT", "BOOL"], {"default": "STRING"}), + "remote_nodeid2": ("STRING", {"default": ""}), + "remote_param2": ("STRING", {"default": ""}), + "remote_value2": ("STRING", {"default": ""}), + "remote_type2": (["STRING", "INT", "FLOAT", "BOOL"], {"default": "STRING"}), + "remote_nodeid3": ("STRING", {"default": ""}), + "remote_param3": ("STRING", {"default": ""}), + "remote_value3": ("STRING", {"default": ""}), + "remote_type3": (["STRING", "INT", "FLOAT", "BOOL"], {"default": "STRING"}), + "remote_nodeid4": ("STRING", {"default": ""}), + "remote_param4": ("STRING", {"default": ""}), + "remote_value4": ("STRING", {"default": ""}), + "remote_type4": (["STRING", "INT", "FLOAT", "BOOL"], {"default": "STRING"}), + } + } + + RETURN_TYPES = ("REMCHAIN",) + RETURN_NAMES = ("remote_chain",) + FUNCTION = "chain_start" + CATEGORY = "remote/advanced" + TITLE = "Queue on remote (start of chain)" + + def chain_start(self, workflow, trigger, batch, seed, + remote_nodeid1="", remote_param1="", remote_value1="", remote_type1="STRING", + remote_nodeid2="", remote_param2="", remote_value2="", remote_type2="STRING", + remote_nodeid3="", remote_param3="", remote_value3="", remote_type3="STRING", + remote_nodeid4="", remote_param4="", remote_value4="", remote_type4="STRING"): + + # Make a deep copy of the workflow to avoid retaining state between calls + workflow = copy.deepcopy(workflow) + + remote_params = {} + for nodeid, param, value, value_type in [ + (remote_nodeid1, remote_param1, remote_value1, remote_type1), + (remote_nodeid2, remote_param2, remote_value2, remote_type2), + (remote_nodeid3, remote_param3, remote_value3, remote_type3), + (remote_nodeid4, remote_param4, remote_value4, remote_type4), + ("", "seed", seed, "INT") + ]: + if param and value: + remote_params[(nodeid, param)] = self.parse_value(value, value_type) + + # Apply remote parameters to the prompt + for (nodeid, param), value in remote_params.items(): + truncated_value = str(value)[:30] + '...' if len(str(value)) > 30 else str(value) # Truncate value for display + print(f"Applying param: {param}, value: {truncated_value}, nodeid: {nodeid}") # Debug statement + if nodeid: + if nodeid in workflow: + if param in workflow[nodeid].get("inputs", {}): + workflow[nodeid]["inputs"][param] = value + print(f"Updated node {nodeid} param {param} with value {truncated_value}") # Debug statement + else: + print(f"Param {param} not found in node {nodeid} inputs") # Debug statement + else: + print(f"Node {nodeid} not found in workflow") # Debug statement + else: + for node_key, node_data in workflow.items(): + if param in node_data.get("inputs", {}): + workflow[node_key]["inputs"][param] = value + print(f"Updated node {node_key} param {param} with value {truncated_value}") # Debug statement + break + else: + print(f"Param {param} not found in node {node_key} inputs") # Debug statement + + remote_chain = { + "seed": seed, + "batch": batch, + "prompt": workflow, + "seed_offset": batch, + "job_id": get_new_job_id(), + } + return(remote_chain,) + + @classmethod + def IS_CHANGED(self, workflow, trigger, batch, seed, prompt): + uuid = f"W:{workflow},B:{batch},S:{seed}" + return uuid if trigger == "on_change" else str(time.time()) + + def parse_value(self, value, value_type): + """Parse the value string to the specified type.""" + if value_type == "INT": + return int(value) + elif value_type == "FLOAT": + return float(value) + elif value_type == "BOOL": + return value.lower() == "true" + else: # STRING or any other type + return value.replace("\\", "\\\\") + +class RemoteChainEnd: + """Split [REMCHAIN] into local seed/batch""" + def __init__(self): + pass + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "remote_chain": ("REMCHAIN",) + } + } + + RETURN_TYPES = ("INT", "INT") + RETURN_NAMES = ("seed", "batch") + FUNCTION = "chain_end" + CATEGORY = "remote/advanced" + TITLE = "Queue on remote (end of chain)" + + def chain_end(self, remote_chain): + seed = remote_chain["seed"] + batch = remote_chain["batch"] + return(seed,batch) + +class RemoteQueueWorker: + """Start job on remote worker""" + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "remote_chain": ("REMCHAIN",), + "remote_url": ("STRING", { + "multiline": False, + "default": "http://127.0.0.1:8288/", + }), + "batch_override": ("INT", {"default": 0, "min": 0, "max": 8}), + "enabled": (["true", "false", "remote"],{"default": "true"}), + "outputs": (["final_image", "any"],{"default":"final_image"}), + }, + } + + RETURN_TYPES = ("REMCHAIN", "REMINFO") + RETURN_NAMES = ("remote_chain", "remote_info") + FUNCTION = "queue" + CATEGORY = "remote/advanced" + TITLE = "Queue on remote (worker)" + + def queue(self, remote_chain, remote_url, batch_override, enabled, outputs, + ): + current_offset = remote_chain["seed_offset"] + remote_chain["seed_offset"] += 1 if batch_override == 0 else batch_override + if enabled == "false": + return (remote_chain, {}) + if enabled == "remote": + # apply offset from previous nodes in chain + remote_chain["seed"] += current_offset + if batch_override > 0: + remote_chain["batch"] = batch_override + return (remote_chain, {}) + + remote_url = clean_url(remote_url) + clear_remote_queue(remote_url) + + # Prepare remote parameters + remote_params = {} + + dispatch_to_remote( + remote_url, + remote_chain["prompt"], + remote_chain["job_id"], + remote_params, + outputs, + ) + remote_info = { + "remote_url" : remote_url, + "job_id" : remote_chain["job_id"], + } + return (remote_chain, remote_info) + +NODE_CLASS_MAPPINGS = { + "RemoteApplyValues(Nux)": RemoteApplyValues, + "RemoteChainStart(Nux)": RemoteChainStartNux, + "RemoteChainStart" : RemoteChainStart, + "RemoteQueueWorker" : RemoteQueueWorker, + "RemoteChainEnd" : RemoteChainEnd, + "RemoteApplyValuesMulti(Nux)": RemoteApplyValuesMulti, +} diff --git a/nodes/images.py b/nodes/images.py new file mode 100644 index 0000000..10dc95a --- /dev/null +++ b/nodes/images.py @@ -0,0 +1,127 @@ +import os +import json +import torch +import requests +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo +from base64 import b64encode +from io import BytesIO + +class LoadImageUrl: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "url": ("STRING", { "multiline": False, }) + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "load_image_url" + CATEGORY = "remote/image" + TITLE = "Load Image (URL)" + + def load_image_url(self, url): + with requests.get(url, stream=True) as r: + r.raise_for_status() + i = Image.open(r.raw) + image = i.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + else: + mask = torch.zeros((64,64), dtype=torch.float32, device="cpu") + return (image, mask) + +class SaveImageUrl: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", ), + "url": ("STRING", { "multiline": False, }), + "filename_prefix": ("STRING", {"default": "ComfyUI"}), + "data_format": (["HTML_image", "Raw_data"],) + }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + RETURN_TYPES = () + OUTPUT_NODE = True + FUNCTION = "save_images" + CATEGORY = "remote/image" + TITLE = "Save Image (URL)" + + def save_images(self, images, url, data_format, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): + filename = os.path.basename(os.path.normpath(filename_prefix)) + + counter = 1 + data = {} + for image in images: + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + meta = PngInfo() + if prompt is not None: + meta.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + meta.add_text(x, json.dumps(extra_pnginfo[x])) + + file = f"{filename}_{counter:05}.png" + + buffer = BytesIO() + img.save(buffer, "png", pnginfo=meta, compress_level=4) + buffer.seek(0) + encoded = b64encode(buffer.read()).decode('utf-8') + data[file] = f"data:image/png;base64,{encoded}" if data_format == "HTML_image" else encoded + counter += 1 + + with requests.post(url, json=data) as r: + r.raise_for_status() + return () + +class CombineImageBatch: + """ + This isn't needed anymore but I used it in too many places so I'm keeping it... + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images_a": ("IMAGE",), + "images_b": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "combine_images" + CATEGORY = "remote/image" + TITLE = "Combine images" + + def combine_images(self,images_a,images_b): + try: + out = torch.cat((images_a,images_b), 0) + except RuntimeError: + print(f"Imagine size mismatch! {images_a.size()}, {images_b.size()}") + out = images_a + return (out,) + + +NODE_CLASS_MAPPINGS = { + "LoadImageUrl" : LoadImageUrl, + "SaveImageUrl" : SaveImageUrl, + "CombineImageBatch" : CombineImageBatch, +} diff --git a/nodes/latents.py b/nodes/latents.py new file mode 100644 index 0000000..33b356d --- /dev/null +++ b/nodes/latents.py @@ -0,0 +1,629 @@ +import os +import torch +import requests +import safetensors.torch +import numpy as np +import io +from io import BytesIO +import json +from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageEnhance +from PIL.PngImagePlugin import PngInfo +import folder_paths +import base64 + + + +class LoadLatentNumpy: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + exts = [".latent", ".safetensors", ".npy", ".npz"] + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + files = [f for f in files if any([f.endswith(x) for x in exts])] + return { + "required": { + "latent": [sorted(files), ] + }, + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "load" + CATEGORY = "remote/latent" + TITLE = "Load Latent (Numpy)" + + def load_comfy(self, file): + # From default node - renamed safetensors file + if type(file) == str: + data = safetensors.torch.load_file(file) + else: + data = safetensors.torch.load(file) + + latent = data["latent_tensor"].to(torch.float32) + if "latent_format_version_0" not in data: + latent *= 1.0 / 0.18215 # XL? + return latent + + def load_numpy(self, file): + # plain npy file - saved as-is + return torch.from_numpy(np.load(file)) + + def load_koyha(self, file): + # generated by sd_scripts - npz + if "latents" in data.keys(): + latent = data["latents"] + else: + latent = [x for x in data.items() if x.shape > 3][0] + return torch.from_numpy(latent) + + def load(self, latent): + path = folder_paths.get_annotated_filepath(latent) + name, ext = os.path.splitext(latent) + + if ext in [".latent", ".safetensors"]: + latent = self.load_comfy(path) + elif ext == ".npy": + latent = self.load_numpy(path) + elif ext == ".npz": + latent = self.load_koyha(path) + else: + try: + latent = self.load_numpy(path) + except: + raise ValueError(f"Unknown latent extension '{ext}'") + + if len(latent.shape) == 3: + latent = latent.unsqueeze(0) + print("asdasd", latent.shape) + + return ({"samples": latent.to(torch.float32)},) + + @classmethod + def IS_CHANGED(s, latent): + image_path = folder_paths.get_annotated_filepath(latent) + m = hashlib.sha256() + with open(image_path, 'rb') as f: + m.update(f.read()) + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(s, latent): + if not folder_paths.exists_annotated_filepath(latent): + return f"Invalid latent file '{latent}'" + return True + +class LoadLatentUrl(LoadLatentNumpy): + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "url": ("STRING", { "multiline": False, }) + } + } + + RETURN_TYPES = ("LATENT",) + TITLE = "Load Latent (URL)" + + def load(self, url): + buffer = BytesIO() + with requests.get(url, stream=True, timeout=16) as r: + r.raise_for_status() + buffer.write(r.content) + buffer.seek(0) + + if ".latent" in url or ".safetensors" in url: + latent = self.load_comfy(buffer) + elif ".npy" in url: + latent = self.load_numpy(buffer) + elif ".npz" in url: + latent = self.load_koyha(buffer) + else: + try: + latent = self.load_comfy(buffer) + except: + raise ValueError(f"Unknown latent extension '{url}'") + + if len(latent.shape) == 3: + latent = latent.unsqueeze(0) + + del buffer + return ({"samples": latent.to(torch.float32)},) + + @classmethod + def IS_CHANGED(s, url): + return str(url) + + @classmethod + def VALIDATE_INPUTS(s, url): + return True + +class SaveLatentNumpy: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "samples": ("LATENT",), + "filename_prefix": ("STRING", {"default": "latents/ComfyUI"}) + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("filename",) + OUTPUT_NODE = True + FUNCTION = "save" + CATEGORY = "remote/latent" + TITLE = "Save Latent (Numpy)" + + def save(self, samples, filename_prefix="ComfyUI"): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) + fname = f"{filename}_{counter:05}_.npy" + path = os.path.join(full_output_folder, fname) + np.save(path, samples["samples"].numpy()) + return (fname,) + + +class LatentToBase64Nux: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "samples": ("LATENT",), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("base64_latent",) + FUNCTION = "convert" + CATEGORY = "remote/latent" + TITLE = "Latent to Base64" + + def convert(self, samples): + # Convert the latent samples to a numpy array + latent_array = samples["samples"].numpy() + + # Save the numpy array to a bytes buffer + buffer = io.BytesIO() + np.save(buffer, latent_array) + buffer.seek(0) + + # Encode the bytes to base64 + base64_latent = base64.b64encode(buffer.getvalue()).decode('utf-8') + + return (base64_latent,) + + +class LoadLatentFromBase64Nux: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "base64_latent": ("STRING", {"multiline": True, "default": "", "forceInput": True}, ), + }, + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "load" + CATEGORY = "remote/latent" + TITLE = "Load Latent from Base64" + + def load(self, base64_latent=""): + try: + # Decode the base64 string + decoded_data = base64.b64decode(base64_latent) + + # Load the numpy array from the decoded data + buffer = io.BytesIO(decoded_data) + latent_array = np.load(buffer) + + # Convert numpy array to torch tensor + latent_tensor = torch.from_numpy(latent_array).to(torch.float32) + + # Ensure the tensor has the correct shape (add batch dimension if necessary) + if len(latent_tensor.shape) == 3: + latent_tensor = latent_tensor.unsqueeze(0) + + print("Loaded latent shape:", latent_tensor.shape) + + return ({"samples": latent_tensor},) + + except Exception as e: + raise ValueError(f"Failed to load latent from base64: {str(e)}") + + @classmethod + def IS_CHANGED(s, base64_latent): + # Since the input is a string, we can use its hash as a change indicator + return hash(base64_latent) + + @classmethod + def VALIDATE_INPUTS(s, base64_latent): + if not base64_latent: + return "Base64 latent string is empty" + try: + decoded_data = base64.b64decode(base64_latent) + buffer = io.BytesIO(decoded_data) + np.load(buffer) + except: + return "Invalid base64 latent string" + return True + + + +class ConditioningToBase64: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "conditioning": ("CONDITIONING", {"tooltip": "The conditioning to be encoded as base64."}), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("base64_conditioning",) + FUNCTION = "convert" + CATEGORY = "conditioning" + TITLE = "Conditioning2Base64" + + def convert(self, conditioning): + # Extract the conditioning data + cond_data, cond_meta = conditioning[0] + + # Convert tensors to numpy arrays + cond_data_np = cond_data.cpu().numpy() + cond_meta_serializable = {k: v.cpu().numpy() if isinstance(v, torch.Tensor) else v for k, v in cond_meta.items()} + + # Combine data and metadata into a single dictionary + combined_data = { + "cond_data": cond_data_np, + "cond_meta": cond_meta_serializable + } + + # Save the combined data to a bytes buffer + buffer = io.BytesIO() + np.savez_compressed(buffer, **combined_data) + buffer.seek(0) + + # Encode the bytes to base64 + base64_conditioning = base64.b64encode(buffer.getvalue()).decode('utf-8') + + return (base64_conditioning,) + +class ConditioningFromBase64: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "base64_conditioning": ("STRING", {"multiline": True, "forceInput": True}), + } + } + + RETURN_TYPES = ("CONDITIONING",) + FUNCTION = "convert" + CATEGORY = "conditioning" + TITLE = "Remote Conditioning" + + def convert(self, base64_conditioning): + try: + # Decode the base64 string + decoded_data = base64.b64decode(base64_conditioning) + + # Load the numpy arrays from the decoded data + buffer = io.BytesIO(decoded_data) + loaded_data = np.load(buffer, allow_pickle=True) + + # Extract the conditioning data and metadata + cond_data_np = loaded_data['cond_data'] + cond_meta_serializable = loaded_data['cond_meta'].item() + + # Convert numpy arrays back to tensors + cond_data = torch.from_numpy(cond_data_np) + cond_meta = {k: torch.from_numpy(v) if isinstance(v, np.ndarray) else v for k, v in cond_meta_serializable.items()} + + # Reconstruct the conditioning object + conditioning = [[cond_data, cond_meta]] + + return (conditioning,) + + except Exception as e: + raise ValueError(f"Failed to load conditioning from base64: {str(e)}") + +def align_text(align, img_height, text_height, text_pos_y, margins): + if align == "center": + text_plot_y = img_height / 2 - text_height / 2 + text_pos_y + elif align == "top": + text_plot_y = text_pos_y + margins + elif align == "bottom": + text_plot_y = img_height - text_height + text_pos_y - margins + return text_plot_y + + +def justify_text(justify, img_width, line_width, margins): + if justify == "left": + text_plot_x = 0 + margins + elif justify == "right": + text_plot_x = img_width - line_width - margins + elif justify == "center": + text_plot_x = img_width/2 - line_width/2 + return text_plot_x + + +def get_text_size(draw, text, font): + bbox = draw.textbbox((0, 0), text, font=font) + + # Calculate the text width and height + text_width = bbox[2] - bbox[0] + text_height = bbox[3] - bbox[1] + return text_width, text_height + +def draw_masked_text(text_mask, text, + font_name, font_size, + margins, line_spacing, + position_x, position_y, + align, justify, + rotation_angle, rotation_options): + + # Create the drawing context + draw = ImageDraw.Draw(text_mask) + + # Define font settings + font_folder = "fonts" + font_file = os.path.join(font_folder, font_name) + resolved_font_path = os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))), font_file) + font = ImageFont.truetype(str(resolved_font_path), size=font_size) + + # Split the input text into lines + text_lines = text.split('\n') + + # Calculate the size of the text plus padding for the tallest line + max_text_width = 0 + max_text_height = 0 + + for line in text_lines: + # Calculate the width and height of the current line + line_width, line_height = get_text_size(draw, line, font) + + line_height = line_height + line_spacing + max_text_width = max(max_text_width, line_width) + max_text_height = max(max_text_height, line_height) + + # Get the image width and height + image_width, image_height = text_mask.size + image_center_x = image_width / 2 + image_center_y = image_height / 2 + + text_pos_y = position_y + sum_text_plot_y = 0 + text_height = max_text_height * len(text_lines) + + for line in text_lines: + # Calculate the width of the current line + line_width, _ = get_text_size(draw, line, font) + + # Get the text x and y positions for each line + text_plot_x = position_x + justify_text(justify, image_width, line_width, margins) + text_plot_y = align_text(align, image_height, text_height, text_pos_y, margins) + + # Add the current line to the text mask + draw.text((text_plot_x, text_plot_y), line, fill=255, font=font) + + text_pos_y += max_text_height # Move down for the next line + sum_text_plot_y += text_plot_y # Sum the y positions + + # Calculate centers for rotation + text_center_x = text_plot_x + max_text_width / 2 + text_center_y = sum_text_plot_y / len(text_lines) + + if rotation_options == "text center": + rotated_text_mask = text_mask.rotate(rotation_angle, center=(text_center_x, text_center_y)) + elif rotation_options == "image center": + rotated_text_mask = text_mask.rotate(rotation_angle, center=(image_center_x, image_center_y)) + + return rotated_text_mask + + +class SaveImageWithBase64: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + self.type = "output" + self.prefix_append = "" + self.compress_level = 4 + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", {"tooltip": "The images to save."}), + "filename_prefix": ("STRING", {"default": "ComfyUI",}) + }, + "optional": { + "workflowName": ("STRING", {"default": "",}), + "latent": ("LATENT",), + "positive_conditioning": ("CONDITIONING",), + "negative_conditioning": ("CONDITIONING",) + }, + "hidden": { + "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" + }, + } + + RETURN_TYPES = () + FUNCTION = "save_images" + OUTPUT_NODE = True + CATEGORY = "image" + TITLE = "save conds and latents" + + def save_images(self, images, filename_prefix="ComfyUI", workflowName="", latent=None, positive_conditioning=None, negative_conditioning=None, prompt=None, extra_pnginfo=None): + + def convertconditioning(conditioning): + cond_data, cond_meta = conditioning[0] + # Convert tensors to numpy arrays + cond_data_np = cond_data.cpu().numpy() + cond_meta_serializable = {k: v.cpu().numpy() if isinstance(v, torch.Tensor) else v for k, v in cond_meta.items()} + # Combine data and metadata into a single dictionary + combined_data = { + "cond_data": cond_data_np, + "cond_meta": cond_meta_serializable + } + # Save the combined data to a bytes buffer + buffer = io.BytesIO() + np.savez_compressed(buffer, **combined_data) + buffer.seek(0) + # Encode the bytes to base64 + base64_conditioning = base64.b64encode(buffer.getvalue()).decode('utf-8') + return base64_conditioning + + def convertlatent(samples): + # Convert the latent samples to a numpy array + latent_array = samples["samples"].numpy() + # Save the numpy array to a bytes buffer + buffer = io.BytesIO() + np.save(buffer, latent_array) + buffer.seek(0) + # Encode the bytes to base64 + base64_latent = base64.b64encode(buffer.getvalue()).decode('utf-8') + return base64_latent + + filename_prefix += self.prefix_append + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) + results = list() + for (batch_number, image) in enumerate(images): + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + + # Add text overlay if workflowName is not empty + if workflowName: + #credits to comfyroll for text overlay + # Convert tensor images + image_3d = images[0, :, :, :] + # Create PIL images for the text and background layers and text mask + back_image = tensor2pil(image_3d) + text_image = Image.new('RGB', back_image.size, (255, 255, 255)) + text_mask = Image.new('L', back_image.size) + + # Draw the text on the text mask + rotated_text_mask = draw_masked_text(text_mask, workflowName, "Roboto-Regular.ttf", 30, + 0, 0, + 0, 0, + 'center', 'center', + 0, 'text center') + + # Composite the text image onto the background image using the rotated text mask + img = Image.composite(text_image, back_image, rotated_text_mask) + + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + if latent is not None: + latent_base64 = convertlatent(latent) + metadata.add_text("latent_base64", latent_base64) + if positive_conditioning is not None: + p_conditioning_base64 = convertconditioning(positive_conditioning) + metadata.add_text("conditioning_base64", p_conditioning_base64) + if negative_conditioning is not None: + n_conditioning_base64 = convertconditioning(negative_conditioning) + metadata.add_text("conditioning_base64", n_conditioning_base64) + filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) + file = f"{filename_with_batch_num}_{counter:05}_.png" + img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + counter += 1 + return { "ui": { "images": results } } + + +class ExtractBase64FromImageUpload: + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + return {"required": + {"image": (sorted(files), {"image_upload": True})}, + } + + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("latent_base64", "conditioning_base64") + FUNCTION = "extract" + CATEGORY = "image" + TITLE = "Extracts base64 encoded latent and conditioning data from an image's metadata." + + def extract(self, image): + image_path = folder_paths.get_annotated_filepath(image) + img = Image.open(image_path) + + latent_base64 = None + conditioning_base64 = None + + if "latent_base64" in img.info: + latent_base64 = img.info["latent_base64"] + + if "conditioning_base64" in img.info: + conditioning_base64 = img.info["conditioning_base64"] + + return (latent_base64, conditioning_base64) + + @classmethod + def IS_CHANGED(s, image): + image_path = folder_paths.get_annotated_filepath(image) + m = hashlib.sha256() + with open(image_path, 'rb') as f: + m.update(f.read()) + return m.digest().hex() + +def tensor2pil(t_image: torch.Tensor) -> Image: + return Image.fromarray(np.clip(255.0 * t_image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +class ExtractBase64FromImage: + @classmethod + def INPUT_TYPES(s): + return {"required": {"image": ("IMAGE",)}} + + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("latent_base64", "conditioning_base64") + FUNCTION = "extract" + CATEGORY = "image" + TITLE = "Extracts base64 encoded latent and conditioning data from an image tensor's metadata." + + def extract(self, image): + # Access the tensor directly + tensor = image[0] + + latent_base64 = None + conditioning_base64 = None + + # Retrieve metadata from the tensor + if hasattr(tensor, 'metadata'): + metadata = tensor.metadata + print("Metadata keys:", list(metadata.keys())) # Print the metadata keys + latent_base64 = metadata.get("latent_base64", None) + conditioning_base64 = metadata.get("conditioning_base64", None) + + return (latent_base64, conditioning_base64) + +NODE_CLASS_MAPPINGS = { + "LoadLatentFromBase64(Nux)": LoadLatentFromBase64Nux, + "LatentToBase64(Nux)": LatentToBase64Nux, + "LoadLatentNumpy": LoadLatentNumpy, + "LoadLatentUrl": LoadLatentUrl, + "SaveLatentNumpy": SaveLatentNumpy, + "ConditioningToBase64(Nux)": ConditioningToBase64, # New class + "ConditioningFromBase64(Nux)": ConditioningFromBase64, + "SaveImageWithBase64(Nux)": SaveImageWithBase64, + "ExtractBase64FromImage(Nux)": ExtractBase64FromImage, + "ExtractBase64FromImageUpload(Nux)": ExtractBase64FromImageUpload +} \ No newline at end of file diff --git a/nodes/simple.py b/nodes/simple.py new file mode 100644 index 0000000..d383d3c --- /dev/null +++ b/nodes/simple.py @@ -0,0 +1,289 @@ +from ..core.fetch import fetch_from_remote, fetch_from_remote_with_extras +from ..core.utils import clean_url, get_client_id, get_new_job_id +from ..core.dispatch import dispatch_to_remote, clear_remote_queue + +class FetchRemote(): + """ + Try to retrieve the final output image from the remote client. + On the remote client, this is replaced with a preview image node. + I.e. remote_info can be none, but the node shouldn't exist at that point + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "final_image": ("IMAGE",), + "remote_info": ("REMINFO",), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "fetch" + CATEGORY = "remote" + TITLE = "Fetch from remote" + + def fetch(self, final_image, remote_info): + out = fetch_from_remote( + remote_url = remote_info.get("remote_url"), + job_id = remote_info.get("job_id"), + ) + if out is None: + out = final_image[:1] * 0.0 # black image + return (out,) + +#with extras returns, the image, remote latent and conditioning if there are any +class FetchRemoteWithExtras(): + """ + Try to retrieve the final output image from the remote client. + On the remote client, this is replaced with a preview image node. + I.e. remote_info can be none, but the node shouldn't exist at that point + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "final_image": ("IMAGE",), + "remote_info": ("REMINFO",), + }, + } + + RETURN_TYPES = ("IMAGE", "STRING", "STRING") + RETURN_NAMES = ("IMAGE", "latent_base64", "conditioning_base64") + FUNCTION = "fetch" + CATEGORY = "remote" + TITLE = "Fetch from remote" + + def fetch(self, final_image, remote_info): + out, metadata = fetch_from_remote_with_extras( + remote_url = remote_info.get("remote_url"), + job_id = remote_info.get("job_id"), + ) + if out is None: + out = final_image[:1] * 0.0 # black image + + latent = metadata.get("latent_base64", None) + conditioning = metadata.get("conditioning_base64", None) + + return (out, latent, conditioning) + +import ast + +#v2 of remote queue to prevent to fix ui mess +class RemoteQueueSimpleNux(): + """ + This is a "simplified" version with additional optional parameters for remote execution. + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "remote_url": ("STRING", { + "multiline": False, + "default": "http://127.0.0.1:8288/", + }), + "batch_local": ("INT", {"default": 1, "min": 1, "max": 8}), + "batch_remote": ("INT", {"default": 1, "min": 1, "max": 8}), + "trigger": (["on_change", "always"],), + "enabled": (["true", "false", "remote"],{"default": "true"}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + }, + "optional": { + "remoteapply1": ("REMOTEAPPLY",), + "remoteapply2": ("REMOTEAPPLY",), + "remoteapply3": ("REMOTEAPPLY",), + "remoteapply4": ("REMOTEAPPLY",), + "remoteapply5": ("REMOTEAPPLY",), + "remoteapply6": ("REMOTEAPPLY",), + "remoteapply7": ("REMOTEAPPLY",), + "remoteapply8": ("REMOTEAPPLY",), + "remoteapply9": ("REMOTEAPPLY",), + "remoteapply10": ("REMOTEAPPLY",), + }, + "hidden": { + "prompt": "PROMPT", + }, + } + + RETURN_TYPES = ("INT", "INT", "REMINFO",) + RETURN_NAMES = ("seed", "batch", "remote_info",) + FUNCTION = "queue" + CATEGORY = "remote" + TITLE = "Queue on remote (single)" + + def parse_value(self, value, value_type): + """Parse the value string to the specified type.""" + if value_type == "INT": + return int(value) + elif value_type == "FLOAT": + return float(value) + elif value_type == "BOOL": + return value.lower() == "true" + else: # STRING or any other type + return value + + def queue(self, remote_url, batch_local, batch_remote, trigger, enabled, seed, prompt, + remoteapply1=None, remoteapply2=None, remoteapply3=None, remoteapply4=None, + remoteapply5=None, remoteapply6=None, remoteapply7=None, remoteapply8=None, + remoteapply9=None, remoteapply10=None): + if enabled == "false": + return (seed, batch_local, {}) + if enabled == "remote": + return (seed+batch_local, batch_remote, {}) + + job_id = get_new_job_id() + remote_url = clean_url(remote_url) + clear_remote_queue(remote_url) + + # Prepare remote parameters + remote_params = [] + for remoteapply in [remoteapply1, remoteapply2, remoteapply3, remoteapply4, + remoteapply5, remoteapply6, remoteapply7, remoteapply8, + remoteapply9, remoteapply10]: + if remoteapply: + nodetitle, param, value, value_type = remoteapply + if param and value: + remote_params.append((param, self.parse_value(value, value_type), nodetitle)) + + dispatch_to_remote(remote_url, prompt, job_id, remote_params) + remote_info = { + "remote_url" : remote_url, + "job_id" : job_id, + } + return (seed, batch_local, remote_info) + + @classmethod + def IS_CHANGED(self, remote_url, batch_local, batch_remote, trigger, enabled, seed, prompt, + remoteapply1=None, remoteapply2=None, remoteapply3=None, remoteapply4=None, + remoteapply5=None, remoteapply6=None, remoteapply7=None, remoteapply8=None, + remoteapply9=None, remoteapply10=None): + uuid = f"W:{remote_url},B1:{batch_local},B2:{batch_remote},S:{seed},E:{enabled}" + for i, remoteapply in enumerate([remoteapply1, remoteapply2, remoteapply3, remoteapply4, + remoteapply5, remoteapply6, remoteapply7, remoteapply8, + remoteapply9, remoteapply10], start=1): + if remoteapply: + param, value, value_type, nodetitle = remoteapply + uuid += f",RP{i}:{param}:{value}:{value_type}:{nodetitle}" + return uuid if trigger == "on_change" else str(time.time()) + +class RemoteQueueSimple(): + """ + This is a "simplified" version with additional optional parameters for remote execution. + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "remote_url": ("STRING", { + "multiline": False, + "default": "http://127.0.0.1:8288/", + }), + "batch_local": ("INT", {"default": 1, "min": 1, "max": 8}), + "batch_remote": ("INT", {"default": 1, "min": 1, "max": 8}), + "trigger": (["on_change", "always"],), + "enabled": (["true", "false", "remote"],{"default": "true"}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + }, + #init idea for entering settings for the remote workflow. + "optional": { + "remote_param1": ("STRING", {"default": ""}), + "remote_value1": ("STRING", {"default": ""}), + "remote_type1": (["STRING", "INT", "FLOAT", "BOOL"], {"default": "STRING"}), + "remote_nodetitle1": ("STRING", {"default": ""}), # Added nodetitle + "remote_param2": ("STRING", {"default": ""}), + "remote_value2": ("STRING", {"default": ""}), + "remote_type2": (["STRING", "INT", "FLOAT", "BOOL"], {"default": "STRING"}), + "remote_nodetitle2": ("STRING", {"default": ""}), # Added nodetitle + "remote_param3": ("STRING", {"default": ""}), + "remote_value3": ("STRING", {"default": ""}), + "remote_type3": (["STRING", "INT", "FLOAT", "BOOL"], {"default": "STRING"}), + "remote_nodetitle3": ("STRING", {"default": ""}), # Added nodetitle + "remote_param4": ("STRING", {"default": ""}), + "remote_value4": ("STRING", {"default": ""}), + "remote_type4": (["STRING", "INT", "FLOAT", "BOOL"], {"default": "STRING"}), + "remote_nodetitle4": ("STRING", {"default": ""}), # Added nodetitle + }, + "hidden": { + "prompt": "PROMPT", + }, + } + + RETURN_TYPES = ("INT", "INT", "REMINFO",) + RETURN_NAMES = ("seed", "batch", "remote_info",) + FUNCTION = "queue" + CATEGORY = "remote" + TITLE = "Queue on remote (single)" + + def parse_value(self, value, value_type): + """Parse the value string to the specified type.""" + if value_type == "INT": + return int(value) + elif value_type == "FLOAT": + return float(value) + elif value_type == "BOOL": + return value.lower() == "true" + else: # STRING or any other type + return value + + def queue(self, remote_url, batch_local, batch_remote, trigger, enabled, seed, prompt, + remote_param1="", remote_value1="", remote_type1="STRING", remote_nodetitle1="", + remote_param2="", remote_value2="", remote_type2="STRING", remote_nodetitle2="", + remote_param3="", remote_value3="", remote_type3="STRING", remote_nodetitle3="", + remote_param4="", remote_value4="", remote_type4="STRING", remote_nodetitle4=""): + if enabled == "false": + return (seed, batch_local, {}) + if enabled == "remote": + return (seed+batch_local, batch_remote, {}) + + job_id = get_new_job_id() + remote_url = clean_url(remote_url) + clear_remote_queue(remote_url) + + # Prepare remote parameters + remote_params = [] + for param, value, value_type, nodetitle in [ + (remote_param1, remote_value1, remote_type1, remote_nodetitle1), + (remote_param2, remote_value2, remote_type2, remote_nodetitle2), + (remote_param3, remote_value3, remote_type3, remote_nodetitle3), + (remote_param4, remote_value4, remote_type4, remote_nodetitle4) + ]: + if param and value: + remote_params.append((param, self.parse_value(value, value_type), nodetitle)) + + dispatch_to_remote(remote_url, prompt, job_id, remote_params) + remote_info = { + "remote_url" : remote_url, + "job_id" : job_id, + } + return (seed, batch_local, remote_info) + + @classmethod + def IS_CHANGED(self, remote_url, batch_local, batch_remote, trigger, enabled, seed, prompt, + remote_param1="", remote_value1="", remote_type1="STRING", remote_nodetitle1="", + remote_param2="", remote_value2="", remote_type2="STRING", remote_nodetitle2="", + remote_param3="", remote_value3="", remote_type3="STRING", remote_nodetitle3="", + remote_param4="", remote_value4="", remote_type4="STRING", remote_nodetitle4=""): + uuid = f"W:{remote_url},B1:{batch_local},B2:{batch_remote},S:{seed},E:{enabled}" + uuid += f",RP1:{remote_param1}:{remote_value1}:{remote_type1}:{remote_nodetitle1}" + uuid += f",RP2:{remote_param2}:{remote_value2}:{remote_type2}:{remote_nodetitle2}" + uuid += f",RP3:{remote_param3}:{remote_value3}:{remote_type3}:{remote_nodetitle3}" + uuid += f",RP4:{remote_param4}:{remote_value4}:{remote_type4}:{remote_nodetitle4}" + return uuid if trigger == "on_change" else str(time.time()) + +NODE_CLASS_MAPPINGS = { + "RemoteQueueSimple(Nux)" : RemoteQueueSimpleNux, + "RemoteQueueSimple" : RemoteQueueSimple, + "FetchRemote" : FetchRemote, + "FetchRemoteWithExtras(Nux)": FetchRemoteWithExtras, +} diff --git a/nodes/workflows.py b/nodes/workflows.py new file mode 100644 index 0000000..3ccb302 --- /dev/null +++ b/nodes/workflows.py @@ -0,0 +1,211 @@ +import os +import json +import hashlib +import folder_paths + +class SaveDiskWorkflowJSON: + """Save workflow to disk""" + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "workflow": ("JSON", ), + "filename_prefix": ("STRING", {"default": "workflow/ComfyUI"}), + } + } + + RETURN_TYPES = () + FUNCTION = "save_workflow" + OUTPUT_NODE = True + CATEGORY = "remote/advanced" + TITLE = "Save workflow (disk)" + + def save_workflow(self, workflow, filename_prefix): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) + + json_path = os.path.join(full_output_folder, f"{filename}_{counter:05}_.json") + with open(json_path, "w") as f: + f.write(json.dumps(workflow, indent=2)) + return {} + +class LoadDiskWorkflowJSON: + """Load workflow JSON from disk""" + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.endswith(".json")] + return { + "required": { + "workflow": [sorted(files),], + } + } + + RETURN_TYPES = ("JSON",) + RETURN_NAMES = ("Workflow JSON",) + FUNCTION = "load_workflow" + CATEGORY = "remote/advanced" + TITLE = "Load workflow (disk)" + + def load_workflow(self, workflow): + json_path = folder_paths.get_annotated_filepath(workflow) + with open(json_path) as f: + data = json.loads(f.read()) + return (data,) + + @classmethod + def IS_CHANGED(s, workflow): + json_path = folder_paths.get_annotated_filepath(workflow) + m = hashlib.sha256() + with open(json_path, 'rb') as f: + m.update(f.read()) + return m.digest().hex() + + @classmethod + def VALIDATE_INPUTS(s, workflow): + if not folder_paths.exists_annotated_filepath(workflow): + return "Invalid JSON file: {}".format(workflow) + json_path = folder_paths.get_annotated_filepath(workflow) + with open(json_path) as f: + try: json.loads(f.read()) + except: + return "Failed to read JSON file: {}".format(workflow) + return True + +class LoadCurrentWorkflowJSON: + """Fetch the current workflow/prompt as an API compatible JSON""" + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": {}, + "hidden": { + "prompt": "PROMPT", + }, + } + + RETURN_TYPES = ("JSON",) + RETURN_NAMES = ("Workflow JSON",) + FUNCTION = "load_workflow" + CATEGORY = "remote/advanced" + TITLE = "Load workflow (current)" + + def load_workflow(self, prompt): + return (prompt,) + + @classmethod + def IS_CHANGED(s, prompt): + return hashlib.sha256(json.dumps(prompt)).digest().hex() + +import os +import json +import hashlib +import folder_paths +from PIL import Image +import piexif + + +class LoadWorkflowJSON: + """Load workflow JSON from disk or image upload""" + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))] + return { + "required": { + "image": (sorted(files), {"image_upload": True}), + } + } + + RETURN_TYPES = ("JSON",) + RETURN_NAMES = ("Workflow JSON",) + FUNCTION = "load_workflow" + CATEGORY = "remote/advanced" + TITLE = "Load workflow (image upload)" + + def load_workflow(self, image): + try: + image_path = folder_paths.get_annotated_filepath(image) + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + + with Image.open(image_path) as img: + # Try to get the EXIF data + exif_data = img.info.get("exif") + if not exif_data: + raise ValueError("No EXIF data found in the image") + + exif_dict = piexif.load(exif_data) + user_comment = exif_dict.get("Exif", {}).get(piexif.ExifIFD.UserComment) + if not user_comment: + raise ValueError("No UserComment found in EXIF data") + + metadata = json.loads(user_comment.decode("utf-8")) + workflow_json = metadata.get("prompt") + if not workflow_json: + raise ValueError("No 'prompt' field found in UserComment data") + + # Parse the JSON string into a Python object + workflow_data = json.loads(workflow_json) + return (workflow_data,) + + except Exception as e: + print(f"Error loading workflow from image: {str(e)}") + return ({},) # Return an empty JSON object in case of error + + @classmethod + def IS_CHANGED(s, image): + try: + image_path = folder_paths.get_annotated_filepath(image) + if not os.path.exists(image_path): + return "FILE_NOT_FOUND" + + m = hashlib.sha256() + with open(image_path, 'rb') as f: + m.update(f.read()) + return m.digest().hex() + except Exception as e: + print(f"Error in IS_CHANGED: {str(e)}") + return "ERROR" + + @classmethod + def VALIDATE_INPUTS(s, image): + try: + image_path = folder_paths.get_annotated_filepath(image) + if not os.path.exists(image_path): + return f"Image file not found: {image_path}" + + with Image.open(image_path) as img: + exif_data = img.info.get("exif") + if not exif_data: + return f"No EXIF data found in image: {image}" + + exif_dict = piexif.load(exif_data) + user_comment = exif_dict.get("Exif", {}).get(piexif.ExifIFD.UserComment) + if not user_comment: + return f"No UserComment found in EXIF data: {image}" + + metadata = json.loads(user_comment.decode("utf-8")) + if "prompt" not in metadata: + return f"No 'prompt' field found in UserComment data: {image}" + + return True + except Exception as e: + return f"Error validating image: {image}. Error: {str(e)}" + +NODE_CLASS_MAPPINGS = { + "SaveDiskWorkflowJSON": SaveDiskWorkflowJSON, + "LoadDiskWorkflowJSON": LoadDiskWorkflowJSON, + "LoadCurrentWorkflowJSON": LoadCurrentWorkflowJSON, + "LoadWorkflowJSON": LoadWorkflowJSON, +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..03ff241 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +requests>=2.28.2