Upload rudimentary distribution script

This commit is contained in:
City
2023-09-03 21:03:52 +02:00
parent 8215a3c6df
commit 64fc4a5c46
4 changed files with 198 additions and 0 deletions
+4
View File
@@ -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]
+28
View File
@@ -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"
+7
View File
@@ -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.
+159
View File
@@ -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()