ig_bot/unsorted/tester.py

80 lines
3 KiB
Python
Raw Normal View History

2024-05-25 09:13:33 +02:00
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)