From 64fc4a5c46a5d7e6407b42f579937360014934a7 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sun, 3 Sep 2023 21:03:52 +0200 Subject: [PATCH] Upload rudimentary distribution script --- .gitignore | 4 + mass-process/job.example.yaml | 28 ++++++ mass-process/readme.md | 7 ++ mass-process/server.py | 159 ++++++++++++++++++++++++++++++++++ 4 files changed, 198 insertions(+) create mode 100644 mass-process/job.example.yaml create mode 100644 mass-process/readme.md create mode 100644 mass-process/server.py diff --git a/.gitignore b/.gitignore index 68bc17f..933c844 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,7 @@ +mass-process/output +mass-process/image +mass-process/*.png + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/mass-process/job.example.yaml b/mass-process/job.example.yaml new file mode 100644 index 0000000..f166512 --- /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..beaa95e --- /dev/null +++ b/mass-process/readme.md @@ -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. diff --git a/mass-process/server.py b/mass-process/server.py new file mode 100644 index 0000000..7865076 --- /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()