api_example.py

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")