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) # Update prompt_text with provided arguments 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 the updated values 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', []) # Check position in queue 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.') # Check if the prompt is currently running 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}") # Get base64 encoded image 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 you want to save the base64 image to a file 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")