reworked favourites logic

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}")

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"]

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():