80 lines
3 KiB
Python
80 lines
3 KiB
Python
|
import requests
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
import os
|
||
|
|
import uuid
|
||
|
|
|
||
|
|
# Constants
|
||
|
|
server_address = "127.0.0.1:8188"
|
||
|
|
client_id = str(uuid.uuid4())
|
||
|
|
result_directory = '/Users/sij/AI/Vision/ComfyUI/output'
|
||
|
|
json_file_paths = ['./workflow_api.json']
|
||
|
|
saved_images = '/Users/sij/AI/Vision/igbot/images'
|
||
|
|
|
||
|
|
def queue_prompt(prompt: dict):
|
||
|
|
response = requests.post(f"http://{server_address}/prompt", json={"prompt": prompt, "client_id": client_id})
|
||
|
|
if response.status_code == 200:
|
||
|
|
return response.json().get('prompt_id')
|
||
|
|
else:
|
||
|
|
raise Exception(f"Failed to queue prompt. Status code: {response.status_code}, Response body: {response.text}")
|
||
|
|
|
||
|
|
def poll_status(prompt_id):
|
||
|
|
"""Poll the job status until it's complete and return the status data."""
|
||
|
|
while True:
|
||
|
|
status_response = requests.get(f"http://{server_address}/history/{prompt_id}")
|
||
|
|
print(f"Status code for prompt {prompt_id}: {status_response.status_code}")
|
||
|
|
if status_response.status_code != 200:
|
||
|
|
raise Exception("Failed to get job status")
|
||
|
|
|
||
|
|
status_data = status_response.json()
|
||
|
|
print(f"Status data: {status_data}")
|
||
|
|
job_data = status_data.get(prompt_id, {})
|
||
|
|
if job_data.get("status", {}).get("completed", False):
|
||
|
|
return job_data
|
||
|
|
time.sleep(5) # Poll every 5 seconds
|
||
|
|
|
||
|
|
def get_image(status_data):
|
||
|
|
"""Extract the filename and subfolder from the status data and read the file."""
|
||
|
|
try:
|
||
|
|
outputs = status_data.get("outputs", {})
|
||
|
|
images_info = outputs.get("9", {}).get("images", [])
|
||
|
|
if not images_info:
|
||
|
|
raise Exception("No images found in the job output.")
|
||
|
|
|
||
|
|
image_info = images_info[0] # Assuming the first image is the target
|
||
|
|
filename = image_info.get("filename")
|
||
|
|
subfolder = image_info.get("subfolder", "") # Default to empty if not present
|
||
|
|
file_path = os.path.join(result_directory, subfolder, filename)
|
||
|
|
|
||
|
|
with open(file_path, 'rb') as file:
|
||
|
|
return file.read()
|
||
|
|
except KeyError as e:
|
||
|
|
raise Exception(f"Failed to extract image information due to missing key: {e}")
|
||
|
|
except FileNotFoundError:
|
||
|
|
raise Exception(f"File {filename} not found at the expected path {file_path}")
|
||
|
|
|
||
|
|
# Process each JSON file
|
||
|
|
for json_file_path in json_file_paths:
|
||
|
|
with open(json_file_path, 'r') as file:
|
||
|
|
json_data = json.load(file)
|
||
|
|
|
||
|
|
try:
|
||
|
|
prompt_id = queue_prompt(json_data)
|
||
|
|
print(f"Prompt ID: {prompt_id}")
|
||
|
|
status_data = poll_status(prompt_id)
|
||
|
|
print(f"Status: {status_data}")
|
||
|
|
image_data = get_image(status_data)
|
||
|
|
filename = f"{saved_images}/{prompt_id}.png"
|
||
|
|
image_path = os.path.join(saved_images, filename)
|
||
|
|
try:
|
||
|
|
with open(image_path, 'wb') as file:
|
||
|
|
file.write(image_data)
|
||
|
|
print(f"Image successfully saved to {image_path}")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"An error occurred while saving the image: {e}")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
print(f"An error occurred: {e}")
|
||
|
|
time.sleep(10)
|
||
|
|
|