import random import websocket import uuid import json import urllib.request import urllib.parse import argparse import time import os import base64 from datetime import datetime from PIL import Image import io import glob as glob_module server_address = "127.0.0.1:8188" use_https = False ws_protocol = "ws" token = "" def queue_prompt(prompt): global use_https, ws_protocol client_id = str(uuid.uuid4()) p = {"prompt": prompt, "client_id": client_id} data = json.dumps(p).encode('utf-8') protocol = "https" if use_https else "http" req = urllib.request.Request("{}://{}/prompt?token={}".format(protocol, server_address, token), data=data) return json.loads(urllib.request.urlopen(req).read()), client_id def get_images(ws, prompt): global use_https, ws_protocol prompt_id, client_id = queue_prompt(prompt) print(f"Prompt ID: {prompt_id['prompt_id']}") output_images = {} current_node = "" ws.connect("{}://{}/ws?clientId={}&token={}".format(ws_protocol, server_address, client_id, token)) while True: out = ws.recv() if isinstance(out, str): message = json.loads(out) if message['type'] == 'executing': data = message.get('data', {}) if isinstance(data, dict) and data.get('prompt_id') == prompt_id['prompt_id']: if data.get('node') is None: break else: current_node = data['node'] print(f"Executing node: {current_node}") elif message['type'] == 'executed': data = message.get('data', {}) if isinstance(data, dict) and data.get('prompt_id') == prompt_id['prompt_id']: node_name = data.get('node') if node_name == '126' and 'images' in data: output_images[node_name] = data['images'] if node_name is None: break else: if current_node == '126': images_output = output_images.get(current_node, []) images_output.append(out[8:]) output_images[current_node] = images_output return output_images def upload_image(image_path): global use_https protocol = "https" if use_https else "http" filename = os.path.basename(image_path) boundary = '----WebKitFormBoundary' + ''.join(['-' for _ in range(16)]) with open(image_path, "rb") as f: image_data = f.read() body = '--{0}\r\nContent-Disposition: form-data; name="image"; filename="{1}"\r\nContent-Type: image/{2}\r\n\r\n'.format(boundary, filename, filename.split('.')[-1].lower()) body = body.encode('utf-8') + image_data + ('\r\n--{0}--\r\n'.format(boundary).encode('utf-8')) req = urllib.request.Request( "{}://{}/upload/image?token={}".format(protocol, server_address, token), data=body, headers={"Content-Type": "multipart/form-data; boundary=" + boundary} ) response = json.loads(urllib.request.urlopen(req).read()) return response["name"] parser = argparse.ArgumentParser() parser.add_argument("--url", type=str, default="http://127.0.0.1:8188", help="ComfyUI server URL (default: http://127.0.0.1:8188)") parser.add_argument("--token", type=str, default=None, help="API token") parser.add_argument("--token-file", type=str, default=None, help="File containing API token (default: token.txt)") parser.add_argument("--seed", type=int, default=None) parser.add_argument("--prompt", type=str, default=None, help="Prompt text") parser.add_argument("--9b", "--9b-model", dest="use_9b", action="store_true", help="Use 9B model instead of 4B") parser.add_argument("--output", type=str, default="output", help="Output directory (default: output)") parser.add_argument("--width", type=int, default=None, help="Output width (256-2048)") parser.add_argument("--height", type=int, default=None, help="Output height (256-2048)") parser.add_argument("--force-res", dest="force_res", action="store_true", help="Force resolution instead of keeping aspect ratio") parser.add_argument("--mp", type=float, default=1.0, help="Megapixels for scaling (default: 1.0, max: 4.0)") parser.add_argument("--runs", type=int, default=1, help="Number of times to run the generation per image") parser.add_argument("images", type=str, nargs="*", default=None, help="Image filenames") args = parser.parse_args() if args.token: token = args.token else: if args.token_file: token_path = args.token_file elif os.path.exists("token.txt"): token_path = "token.txt" else: print("Error: No token provided. Create token.txt or use --token or --token-file option.") exit(1) if not os.path.exists(token_path): print(f"Error: Token file '{token_path}' not found.") exit(1) with open(token_path) as f: token = f.read().strip() server_address = args.url use_https = server_address.startswith("https://") if use_https: server_address = server_address.replace("https://", "") ws_protocol = "wss" else: server_address = server_address.replace("http://", "") if server_address.startswith("http://") else server_address ws_protocol = "ws" with open("Restore_Colorize.json") as f: prompt = json.load(f) if args.images: image_files = [] for img_arg in args.images: if any(c in img_arg for c in ['*', '?', '[']): image_files.extend(sorted(glob_module.glob(img_arg))) else: image_files.append(img_arg) else: print("Error: No images provided. Please specify image file(s) as arguments.") exit(1) if not image_files: print("Error: No image files found.") exit(1) print(f"Processing {len(image_files)} image(s), {args.runs} run(s) each") os.makedirs(args.output, exist_ok=True) total_start = time.time() total_images = 0 for idx, img_file in enumerate(image_files): for run_idx in range(args.runs): print(f"\n--- Image {idx+1}/{len(image_files)}, Run {run_idx + 1}/{args.runs}: {img_file} ---") client_id = str(uuid.uuid4()) prompt_copy = json.loads(json.dumps(prompt)) uploaded_name = upload_image(img_file) prompt_copy["76"]["inputs"]["image"] = uploaded_name if "81" in prompt_copy: del prompt_copy["81"] if args.seed is not None: seed = args.seed + run_idx else: seed = random.getrandbits(64) prompt_copy["75:73"]["inputs"]["noise_seed"] = seed if args.use_9b: prompt_copy["75:70"]["inputs"]["unet_name"] = "flux-2-klein-9b-fp8.safetensors" prompt_copy["75:71"]["inputs"]["clip_name"] = "qwen_3_8b_fp8mixed.safetensors" if args.prompt: prompt_copy["75:74"]["inputs"]["text"] = args.prompt if args.force_res: prompt_copy["75:121"]["inputs"]["value"] = True w = args.width if args.width is not None else 1024 h = args.height if args.height is not None else 1024 prompt_copy["75:122"]["inputs"]["value"] = min(max(w, 256), 2048) prompt_copy["75:123"]["inputs"]["value"] = min(max(h, 256), 2048) print(f"Output size: {prompt_copy['75:122']['inputs']['value']}x{prompt_copy['75:123']['inputs']['value']} (forced)") mp = min(max(args.mp, 0.1), 4.0) prompt_copy["75:80"]["inputs"]["megapixels"] = mp if mp != 1.0: print(f"Scaling to {mp} MP") ws = websocket.WebSocket() images = get_images(ws, prompt_copy) ws.close() timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") base_name = os.path.splitext(os.path.basename(img_file))[0] for i, node_id in enumerate(images): for j, image_data in enumerate(images[node_id]): if isinstance(image_data, str): image_data = base64.b64decode(image_data) image = Image.open(io.BytesIO(image_data)) image.save(f"{args.output}/{base_name}_{timestamp}_{seed}_{idx}_{run_idx}_{i}_{j}.png") print(f"Saved: {args.output}/{base_name}_{timestamp}_{seed}_{idx}_{run_idx}_{i}_{j}.png") total_images += 1 elapsed = time.time() - total_start print(f"\nTotal runs completed: {len(image_files) * args.runs}") print(f"Total images: {total_images}") print(f"Total execution time: {elapsed:.2f} seconds")