diff --git a/libs/comfyui.py b/libs/comfyui.py index 9c05052..e4c6d4b 100644 --- a/libs/comfyui.py +++ b/libs/comfyui.py @@ -103,8 +103,7 @@ def generate_image( # Generate image logging.debug(f"Generating image: {file_name}") results = api.queue_and_wait_images(wf, save_node) - rename_image() - + for _, image_data in results.items(): output_path = os.path.join( user_config["comfyui"]["output_dir"], f"{file_name}.png" @@ -112,6 +111,7 @@ def generate_image( with open(output_path, "wb+") as f: f.write(image_data) generate_thumbnail(output_path) + rename_image(f"{file_name}.png") logging.debug(f"Image generated successfully for UID: {file_name}") @@ -143,7 +143,7 @@ def select_model(model: str) -> tuple[str, str]: def create_image(prompt: str | None = None, model: str = "Random") -> None: """Generate an image with a chosen workflow (Random, FLUX*, or SDXL*).""" - + from datetime import datetime if prompt is None: prompt = create_prompt_on_openwebui(user_config["comfyui"]["prompt"]) @@ -153,10 +153,14 @@ def create_image(prompt: str | None = None, model: str = "Random") -> None: save_prompt(prompt) selected_workflow, model = select_model(model) + + # Generate a unique filename using a timestamp + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + file_name = f"{timestamp}_{random.getrandbits(32)}" if selected_workflow == "FLUX": generate_image( - file_name="image", + file_name=file_name, comfy_prompt=prompt, workflow_path="./workflow_flux.json", prompt_node="Positive Prompt T5", @@ -169,6 +173,6 @@ def create_image(prompt: str | None = None, model: str = "Random") -> None: model=model ) else: # SDXL - generate_image("image", comfy_prompt=prompt, model=model) + generate_image(file_name, comfy_prompt=prompt, model=model) logging.info(f"{selected_workflow} generation started with prompt: {prompt}") diff --git a/libs/generic.py b/libs/generic.py index fed2e33..3c24ee8 100644 --- a/libs/generic.py +++ b/libs/generic.py @@ -48,21 +48,41 @@ def load_config() -> configparser.ConfigParser: sys.exit(1) -def rename_image() -> str | None: - """Renames 'image.png' in the output folder to a timestamped filename if it exists.""" - old_path = os.path.join(user_config["comfyui"]["output_dir"], "image.png") +def rename_image(new_filename: str) -> str | None: + """Renames the latest image in the output folder to a timestamped filename.""" + output_dir = user_config["comfyui"]["output_dir"] + image_png_path = os.path.join(output_dir, "image.png") + + # Check if image.png exists and is a favourite + if os.path.exists(image_png_path): + favourites = get_favourites() + if "image.png" in favourites: + # It's a favourite, so rename it to preserve it + timestamped_filename = f"{datetime.now().strftime('%Y%m%d-%H%M%S')}.png" + timestamped_path = os.path.join(output_dir, timestamped_filename) + os.rename(image_png_path, timestamped_path) + + # Update favourites list + favourites.remove("image.png") + favourites.append(timestamped_filename) + save_favourites(favourites) + print(f"Preserved favourite 'image.png' as '{timestamped_filename}'") - 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) - generate_thumbnail(new_path) - print(f"Renamed 'image.png' to '{new_filename}'") - return new_filename - else: - print("No image.png found.") + # Find the latest generated image (which is the new_filename) + latest_image_path = os.path.join(output_dir, new_filename) + if not os.path.exists(latest_image_path): + print(f"Error: Newly generated image '{new_filename}' not found.") return None + # Rename the latest image to "image.png" + if os.path.exists(image_png_path): + os.remove(image_png_path) # remove if it wasn't a favourite + + os.rename(latest_image_path, image_png_path) + generate_thumbnail(image_png_path) + print(f"Renamed '{new_filename}' to 'image.png'") + return "image.png" + def get_details_from_png(path): try: @@ -112,5 +132,18 @@ def load_topics_from_config(): sorted_topics = sorted(topics, key=str.lower) return sorted_topics + +favourites_file = "./favourites.json" + +def get_favourites(): + if not os.path.exists(favourites_file): + return [] + with open(favourites_file, 'r') as f: + return json.load(f) + +def save_favourites(favourites): + with open(favourites_file, 'w') as f: + json.dump(favourites, f) + user_config = load_config() output_folder = user_config["comfyui"]["output_dir"] \ No newline at end of file diff --git a/routes/gallery_routes.py b/routes/gallery_routes.py index 84c1017..a121c31 100644 --- a/routes/gallery_routes.py +++ b/routes/gallery_routes.py @@ -4,17 +4,7 @@ import json bp = Blueprint("gallery_routes", __name__) image_folder = "./output" -favourites_file = "./favourites.json" - -def get_favourites(): - if not os.path.exists(favourites_file): - return [] - with open(favourites_file, 'r') as f: - return json.load(f) - -def save_favourites(favourites): - with open(favourites_file, 'w') as f: - json.dump(favourites, f) +from libs.generic import get_favourites, save_favourites @bp.route("/images", methods=["GET"]) def gallery():