Upload rudimentary distribution script
This commit is contained in:
City
@@ -1,3 +1,7 @@
|
||||
mass-process/output
|
||||
mass-process/image
|
||||
mass-process/*.png
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.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