File size: 2,845 Bytes
42cbd13 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """Render one API graph on your ComfyUI server and download its result."""
import argparse
import json
from pathlib import Path
import time
import urllib.parse
import urllib.request
import uuid
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('workflow', type=Path)
parser.add_argument('--server', default='http://127.0.0.1:8188')
parser.add_argument('--output', required=True, type=Path)
parser.add_argument('--output-node', help='SaveVideo node ID to download for multi-output graphs')
args = parser.parse_args()
if args.output.exists():
parser.error('output already exists; choose a new filename')
graph = json.loads(args.workflow.read_text())
for node in graph.values():
if node['class_type'] == 'SaveVideo':
node['inputs']['filename_prefix'] += '_' + uuid.uuid4().hex[:12]
def api(path, payload=None):
data = None if payload is None else json.dumps(payload).encode()
req = urllib.request.Request(args.server.rstrip('/') + path, data=data,
headers={'Content-Type': 'application/json'})
with urllib.request.urlopen(req, timeout=60) as response:
return json.load(response)
result = api('/prompt', {'prompt': graph, 'client_id': 'fight_quad_example'})
if result.get('node_errors') or 'prompt_id' not in result:
raise RuntimeError(result)
pid = result['prompt_id']
print('Queued', pid, flush=True)
deadline = time.monotonic() + 7200
while time.monotonic() < deadline:
history = api('/history/' + pid)
if pid not in history:
time.sleep(5)
continue
record = history[pid]
if record['status'].get('status_str') != 'success':
raise RuntimeError(record['status'])
for key, node in graph.items():
if node['class_type'] != 'SaveVideo' or (args.output_node and key != args.output_node):
continue
files = record['outputs'].get(key, {}).get('images', [])
for item in files:
if item.get('type') != 'output':
continue
query = urllib.parse.urlencode({k: item[k] for k in ('filename', 'subfolder', 'type')})
with urllib.request.urlopen(args.server.rstrip('/') + '/view?' + query, timeout=120) as response:
with args.output.open('xb') as output:
while chunk := response.read(1024 * 1024):
output.write(chunk)
print('Saved', args.output)
return
raise RuntimeError('Completed without a SaveVideo output')
raise TimeoutError('Render still pending after two hours; check ComfyUI before resubmitting')
if __name__ == '__main__':
main()
|