Upload rudimentary distribution script
This commit is contained in:
@@ -1,3 +1,7 @@
|
|||||||
|
mass-process/output
|
||||||
|
mass-process/image
|
||||||
|
mass-process/*.png
|
||||||
|
|
||||||
# Byte-compiled / optimized / DLL files
|
# Byte-compiled / optimized / DLL files
|
||||||
__pycache__/
|
__pycache__/
|
||||||
*.py[cod]
|
*.py[cod]
|
||||||
|
|||||||
@@ -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"
|
||||||
@@ -0,0 +1,7 @@
|
|||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
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.
|
||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user