| import argparse |
| import json |
| import random |
| import time |
| import requests |
| import base64 |
| from io import BytesIO |
|
|
| def get_image_as_base64(url): |
| try: |
| response = requests.get(url) |
| response.raise_for_status() |
| image_data = BytesIO(response.content) |
| base64_image = base64.b64encode(image_data.getvalue()).decode('utf-8') |
| return base64_image |
| except requests.exceptions.RequestException as ex: |
| print(f'Failed to retrieve image: {ex}') |
| return None |
| |
| def queue_prompt(url, prompt): |
| p = {"prompt": prompt} |
| data = json.dumps(p).encode('utf-8') |
| prompt_url = f"{url}/prompt" |
| try: |
| r = requests.post(prompt_url, data=data) |
| r.raise_for_status() |
| return r.json() |
| except requests.exceptions.RequestException as ex: |
| print(f'POST {prompt_url} failed: {ex}') |
| return None |
|
|
| def get_queue(url): |
| queue_url = f"{url}/queue" |
| try: |
| r = requests.get(queue_url) |
| r.raise_for_status() |
| return r.json() |
| except requests.exceptions.RequestException as ex: |
| print(f'GET {queue_url} failed: {ex}') |
| return None |
|
|
|
|
| def get_history(url, prompt_id): |
| history_url = f"{url}/history/{prompt_id}" |
| try: |
| r = requests.get(history_url) |
| r.raise_for_status() |
| return r.json() |
| except requests.exceptions.RequestException as ex: |
| print(f'GET {history_url} failed: {ex}') |
| return None |
|
|
|
|
| def main(ip, port, filepath, prompt=None, steps=None, seed=None, cfg=None, width=None, height=None, lora_name=None, lora_scale=None): |
| url = f"http://{ip}:{port}" |
|
|
| with open(filepath, 'r') as file: |
| prompt_text = json.load(file) |
|
|
| |
| if prompt is not None: |
| prompt_text["6"]["inputs"]["text"] = prompt |
| if steps is not None: |
| prompt_text["17"]["inputs"]["steps"] = steps |
| if seed is not None: |
| prompt_text["25"]["inputs"]["noise_seed"] = seed |
| else: |
| prompt_text["25"]["inputs"]["noise_seed"] = random.randint(0, 1000000000000000) |
| if cfg is not None: |
| prompt_text["26"]["inputs"]["guidance"] = cfg |
| if width is not None: |
| prompt_text["27"]["inputs"]["width"] = width |
| if height is not None: |
| prompt_text["27"]["inputs"]["height"] = height |
| if lora_name is not None: |
| prompt_text["30"]["inputs"]["lora_name"] = lora_name |
| if lora_scale is not None: |
| prompt_text["30"]["inputs"]["strength_model"] = lora_scale |
|
|
| |
| print(f'Prompt: {prompt_text["6"]["inputs"]["text"]}') |
| print(f'Steps: {prompt_text["17"]["inputs"]["steps"]}') |
| print(f'Seed: {prompt_text["25"]["inputs"]["noise_seed"]}') |
| print(f'CFG: {prompt_text["26"]["inputs"]["guidance"]}') |
| print(f'Width: {prompt_text["27"]["inputs"]["width"]}') |
| print(f'Height: {prompt_text["27"]["inputs"]["height"]}') |
| print(f'LoRA Name: {prompt_text["30"]["inputs"]["lora_name"]}') |
| print(f'LoRA Scale: {prompt_text["30"]["inputs"]["strength_model"]}') |
|
|
| response1 = queue_prompt(url, prompt_text) |
| if response1 is None: |
| print("Failed to queue the prompt.") |
| return |
|
|
| prompt_id = response1['prompt_id'] |
| print(f'Prompt ID: {prompt_id}') |
| print('-' * 20) |
|
|
| while True: |
| time.sleep(5) |
| queue_response = get_queue(url) |
| if queue_response is None: |
| continue |
|
|
| queue_pending = queue_response.get('queue_pending', []) |
| queue_running = queue_response.get('queue_running', []) |
|
|
| |
| for position, item in enumerate(queue_pending): |
| if item[1] == prompt_id: |
| print(f'Queue running: {len(queue_running)}, Queue pending: {len(queue_pending)}, Workflow is in position {position + 1} in the queue.') |
|
|
| |
| for item in queue_running: |
| if item[1] == prompt_id: |
| print(f'Queue running: {len(queue_running)}, Queue pending: {len(queue_pending)}, Workflow is currently running.') |
| break |
|
|
| if not any(prompt_id in item for item in queue_pending + queue_running): |
| break |
|
|
| history_response = get_history(url, prompt_id) |
| if history_response is None: |
| print("Failed to retrieve history.") |
| return |
|
|
| output_info = history_response.get(prompt_id, {}).get('outputs', {}).get('9', {}).get('images', [{}])[0] |
| filename = output_info.get('filename', 'unknown.png') |
| output_url = f"{url}/output/{filename}" |
|
|
| print(f"Output URL: {output_url}") |
|
|
| |
| base64_image = get_image_as_base64(output_url) |
| if base64_image: |
| print("Base64 encoded image:") |
| print(base64_image) |
| else: |
| print("Failed to retrieve base64 encoded image.") |
|
|
| return { |
| "output_url": output_url, |
| "base64_image": base64_image |
| } |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description='Add a prompt to the queue and wait for the output.') |
| parser.add_argument('--ip', type=str, required=True, help='The public IP address of the pod') |
| parser.add_argument('--port', type=int, required=True, help='The external port of the pod') |
| parser.add_argument('--filepath', type=str, required=True, help='The path to the JSON file containing the workflow in api format') |
| parser.add_argument('--prompt', type=str, help='The prompt to use for the workflow') |
| parser.add_argument('--steps', type=int, help='Number of steps for the sampler') |
| parser.add_argument('--seed', type=int, help='Seed for the noise generator') |
| parser.add_argument('--cfg', type=float, help='Classifier-free guidance scale') |
| parser.add_argument('--width', type=int, help='Width of the output image') |
| parser.add_argument('--height', type=int, help='Height of the output image') |
| parser.add_argument('--lora_name', type=str, help='Name of the LoRA to use') |
| parser.add_argument('--lora_scale', type=float, help='Scale of the LoRA effect') |
|
|
| args = parser.parse_args() |
| result = main(args.ip, args.port, args.filepath, args.prompt, args.steps, args.seed, args.cfg, args.width, args.height, args.lora_name, args.lora_scale) |
| |
| |
| if result and result["base64_image"]: |
| with open("output_image.txt", "w") as f: |
| f.write(result["base64_image"]) |
| print("Base64 image saved to output_image.txt") |
