mirror of
https://github.com/karl0ss/ai_image_frame_server.git
synced 2025-04-28 11:13:41 +01:00
91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
import random
|
|
import configparser
|
|
import logging
|
|
import sys
|
|
import litellm
|
|
import time
|
|
import os
|
|
|
|
from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper
|
|
|
|
def load_config():
|
|
user_config = configparser.ConfigParser()
|
|
try:
|
|
user_config.read("./user_config.cfg")
|
|
logging.debug("Configuration loaded successfully.")
|
|
return user_config
|
|
except KeyError as e:
|
|
logging.error(f"Missing configuration key: {e}")
|
|
sys.exit(1)
|
|
|
|
def rename_image():
|
|
"""Rename 'image.png' to a timestamped filename if it exists in the output folder."""
|
|
old_path = os.path.join(user_config["comfyui"]["output_dir"], "image.png")
|
|
|
|
if os.path.exists(old_path):
|
|
new_filename = f"{str(time.time())}.png"
|
|
new_path = os.path.join(user_config["comfyui"]["output_dir"], new_filename)
|
|
os.rename(old_path, new_path)
|
|
print(f"Renamed 'image.png' to '{new_filename}'")
|
|
return new_filename
|
|
else:
|
|
print("No image.png found.")
|
|
return None
|
|
|
|
def send_prompt_to_openwebui(prompt):
|
|
response = litellm.completion(
|
|
api_base=user_config["openwebui"]["base_url"],
|
|
model="openai/" + user_config["openwebui"]["model"],
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": prompt,
|
|
}
|
|
],
|
|
api_key=user_config["openwebui"]["api_key"],
|
|
)
|
|
|
|
return response["choices"][0]["message"]["content"].strip('"')
|
|
|
|
|
|
def generate_image(file_name, comfy_prompt):
|
|
"""Generate an image using the Comfy API."""
|
|
try:
|
|
# Initialize API and workflow
|
|
api = ComfyApiWrapper(user_config["comfyui"]["comfyui_url"])
|
|
wf = ComfyWorkflowWrapper("./workflow_api.json")
|
|
|
|
# Set workflow parameters
|
|
wf.set_node_param("KSampler", "seed", random.getrandbits(32))
|
|
# wf.set_node_param("KSampler", "steps", steps)
|
|
wf.set_node_param("CLIP Text Encode (Prompt)", "text", comfy_prompt)
|
|
wf.set_node_param("Save Image", "filename_prefix", file_name)
|
|
wf.set_node_param("Empty Latent Image", "width", user_config["comfyui"]["width"])
|
|
wf.set_node_param("Empty Latent Image", "height", user_config["comfyui"]["height"])
|
|
wf.set_node_param(
|
|
"Load Checkpoint", "ckpt_name", user_config["comfyui"]["model"]
|
|
)
|
|
# Queue your workflow for completion
|
|
logging.debug(f"Generating image: {file_name}")
|
|
results = api.queue_and_wait_images(wf, "Save Image")
|
|
rename_image()
|
|
for filename, image_data in results.items():
|
|
with open(
|
|
user_config["comfyui"]["output_dir"] + file_name + ".png", "wb+"
|
|
) as f:
|
|
f.write(image_data)
|
|
logging.debug(f"Image generated successfully for UID: {file_name}")
|
|
except Exception as e:
|
|
logging.error(f"Failed to generate image for UID: {file_name}. Error: {e}")
|
|
|
|
|
|
def create_image():
|
|
"""Main function for generating images."""
|
|
prompt = send_prompt_to_openwebui(user_config["comfyui"]["prompt"])
|
|
print(f"Generated prompt: {prompt}")
|
|
generate_image("image", prompt)
|
|
|
|
|
|
user_config = load_config()
|
|
output_folder = user_config["comfyui"]["output_dir"]
|