Karl/ai-frame-image-server
1
0
Fork 0
mirror of https://github.com/karl0ss/ai_image_frame_server.git synced 2025-04-28 19:23:41 +01:00
Code Issues Packages Projects Releases Wiki Activity

rework to single flow

Browse Source
This commit is contained in:
Karl Hudgell 2025-04-19 17:46:04 +01:00
parent a180a7bd4b
commit 30d25d1d7a
1 changed files with 71 additions and 102 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

163
lib.py
View File

@ -6,6 +6,7 @@ import litellm
import time import time
import os import os
import requests import requests
from typing import Optional
from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper
from tenacity import ( from tenacity import (
retry, retry,
@ -48,6 +49,7 @@ def save_prompt(prompt):
f.write(json.dumps(entry) + "\n") f.write(json.dumps(entry) + "\n")
def get_available_models() -> list: def get_available_models() -> list:
"""Fetches available models from ComfyUI.""" """Fetches available models from ComfyUI."""
url = user_config["comfyui"]["comfyui_url"] + "/object_info" url = user_config["comfyui"]["comfyui_url"] + "/object_info"
@ -151,124 +153,80 @@ def create_prompt_on_openwebui(prompt: str) -> str:
# Define the retry logic using Tenacity # Define the retry logic using Tenacity
# @retry( @retry(
# stop=stop_after_attempt(3), stop=stop_after_attempt(3),
# wait=wait_fixed(5), wait=wait_fixed(5),
# before=before_log(logging.getLogger(), logging.DEBUG), before=before_log(logging.getLogger(), logging.DEBUG),
# retry=retry_if_exception_type(Exception) retry=retry_if_exception_type(Exception),
# ) )
def generate_image(file_name: str, comfy_prompt: str) -> None: def generate_image(
"""Generates an image using the Comfy API with retry logic.""" file_name: str,
comfy_prompt: str,
workflow_path: str = "./workflow_api.json",
prompt_node: str = "CLIP Text Encode (Prompt)",
seed_node: str = "KSampler",
seed_param: str = "seed",
save_node: str = "Save Image",
save_param: str = "filename_prefix",
model_node: Optional[str] = "Load Checkpoint",
model_param: Optional[str] = "ckpt_name",
) -> None:
"""Generates an image using the Comfy API with configurable workflow settings."""
try: try:
# Initialize ComfyUI API and workflow
api = ComfyApiWrapper(user_config["comfyui"]["comfyui_url"]) api = ComfyApiWrapper(user_config["comfyui"]["comfyui_url"])
wf = ComfyWorkflowWrapper("./workflow_api.json") wf = ComfyWorkflowWrapper(workflow_path)
# Set workflow parameters # Set workflow parameters
wf.set_node_param(seed_node, seed_param, random.getrandbits(32))
wf.set_node_param(prompt_node, "text", comfy_prompt)
wf.set_node_param(save_node, save_param, file_name)
wf.set_node_param( wf.set_node_param(
"KSampler", "seed", random.getrandbits(32) (
) # Set a random seed for the sampler "Empty Latent Image"
wf.set_node_param( if workflow_path.endswith("workflow_api.json")
"CLIP Text Encode (Prompt)", "text", comfy_prompt else "CR Aspect Ratio"
) # Set the prompt to be used for image generation ),
wf.set_node_param( "width",
"Save Image", "filename_prefix", file_name user_config["comfyui"]["width"],
) # Set the filename prefix for the generated image
wf.set_node_param( # Set image dimensions
"Empty Latent Image", "width", user_config["comfyui"]["width"]
) )
wf.set_node_param( wf.set_node_param(
"Empty Latent Image", "height", user_config["comfyui"]["height"] (
"Empty Latent Image"
if workflow_path.endswith("workflow_api.json")
else "CR Aspect Ratio"
),
"height",
user_config["comfyui"]["height"],
) )
# Validate available models and choose a random one # Conditionally set model if node and param are provided
if model_node and model_param:
valid_models = list( valid_models = list(
set(get_available_models()) # Get all available models from ComfyUI set(get_available_models())
& set(user_config["comfyui"]["models"].split(",")) & set(user_config["comfyui"]["models"].split(","))
) )
if not valid_models: if not valid_models:
raise Exception("No valid options available.") raise Exception("No valid models available.")
model = random.choice(valid_models) model = random.choice(valid_models)
wf.set_node_param( wf.set_node_param(model_node, model_param, model)
"Load Checkpoint", "ckpt_name", model
) # Set the model to be used for image generation
# Generate the image using the workflow and wait for completion # Generate image
logging.debug(f"Generating image: {file_name}") logging.debug(f"Generating image: {file_name}")
results = api.queue_and_wait_images( results = api.queue_and_wait_images(wf, save_node)
wf, "Save Image" rename_image()
) # Queue the workflow and wait for image generation to complete
rename_image() # Rename the generated image file if it exists
# Save the generated image to disk for _, image_data in results.items():
for filename, image_data in results.items(): output_path = os.path.join(
with open( user_config["comfyui"]["output_dir"], f"{file_name}.png"
user_config["comfyui"]["output_dir"] + file_name + ".png", "wb+" )
) as f: with open(output_path, "wb+") as f:
f.write(image_data) f.write(image_data)
logging.debug(f"Image generated successfully for UID: {file_name}") logging.debug(f"Image generated successfully for UID: {file_name}")
except Exception as e: except Exception as e:
logging.error(f"Failed to generate image for UID: {file_name}. Error: {e}") logging.error(f"Failed to generate image for UID: {file_name}. Error: {e}")
raise # Re-raise the exception for Tenacity to handle retries raise
def generate_image_flux(file_name: str, comfy_prompt: str) -> None:
"""Generates an image using the Comfy API with retry logic."""
try:
# Initialize ComfyUI API and workflow
api = ComfyApiWrapper(user_config["comfyui"]["comfyui_url"])
wf = ComfyWorkflowWrapper("./FLUX.json")
# Set workflow parameters
wf.set_node_param(
"Seed", "seed", random.getrandbits(32)
) # Set a random seed for the sampler
wf.set_node_param(
"Positive Prompt T5", "text", comfy_prompt
) # Set the prompt to be used for image generation
wf.set_node_param(
"CivitAI Image Saver", "filename", file_name
) # Set the filename prefix for the generated image
wf.set_node_param( # Set image dimensions
"CR Aspect Ratio", "width", user_config["comfyui"]["width"]
)
wf.set_node_param(
"CR Aspect Ratio", "height", user_config["comfyui"]["height"]
)
# # Validate available models and choose a random one
# valid_models = list(
# set(get_available_models()) # Get all available models from ComfyUI
# & set(user_config["comfyui"]["models"].split(","))
# )
# if not valid_models:
# raise Exception("No valid options available.")
# model = random.choice(valid_models)
# wf.set_node_param(
# "Load Checkpoint", "ckpt_name", model
# ) # Set the model to be used for image generation
# Generate the image using the workflow and wait for completion
logging.debug(f"Generating image: {file_name}")
results = api.queue_and_wait_images(
# wf, "Save Image"
wf,
"CivitAI Image Saver",
) # Queue the workflow and wait for image generation to complete
rename_image() # Rename the generated image file if it exists
# Save the generated image to disk
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}")
raise # Re-raise the exception for Tenacity to handle retries
def create_image(prompt: str | None = None) -> None: def create_image(prompt: str | None = None) -> None:
@ -279,7 +237,18 @@ def create_image(prompt: str | None = None) -> None:
logging.info(f"Generated prompt: {prompt}") # Log generated prompt logging.info(f"Generated prompt: {prompt}") # Log generated prompt
save_prompt(prompt) save_prompt(prompt)
if user_config["comfyui"]["FLUX"]: if user_config["comfyui"]["FLUX"]:
generate_image_flux("image", prompt) generate_image(
file_name="image",
comfy_prompt=prompt,
workflow_path="./FLUX.json",
prompt_node="Positive Prompt T5",
seed_node="Seed",
seed_param="seed",
save_node="CivitAI Image Saver",
save_param="filename",
model_node=None, # FLUX doesn't use model selection
model_param=None,
)
else: else:
generate_image("image", prompt) generate_image("image", prompt)
print(f"Image generation started with prompt: {prompt}") print(f"Image generation started with prompt: {prompt}")