From fc0b1f31664c9d1439dd4e35723523f285a299b6 Mon Sep 17 00:00:00 2001
From: Karl Hudgell <karl.hudgell@bjss.com>
Date: Sat, 29 Mar 2025 08:28:13 +0000
Subject: [PATCH] update with docker server

---
 .gitignore               |  1 -
 Dockerfile               | 17 +++++++++++++++++
 ai_frame_image_server.py | 16 ++++------------
 lib.py                   | 24 +++++++++++++++++-------
 4 files changed, 38 insertions(+), 20 deletions(-)
 create mode 100644 Dockerfile

diff --git a/.gitignore b/.gitignore
index 10f528a..1cd7129 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,5 +5,4 @@ script.log
 build/
 dist/
 user_config.cfg
-Dockerfile
 output/**.*
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000..dc7d0f4
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,17 @@
+# Use an official Python image as a base
+FROM python:3.11
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy project files into the container
+COPY . /app
+
+# Install dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Expose any necessary ports (optional, if running a web app)
+EXPOSE 5000
+
+# Define the command to run the application
+CMD ["python", "ai_frame_image_server.py"]
diff --git a/ai_frame_image_server.py b/ai_frame_image_server.py
index c267f83..cac3535 100644
--- a/ai_frame_image_server.py
+++ b/ai_frame_image_server.py
@@ -6,19 +6,10 @@ app = Flask(__name__)
 
 image_folder = "./output"
 
-def get_latest_image():
-    """Get the latest image file from the directory."""
-    files = [f for f in os.listdir(image_folder) if f.endswith(('.png', '.jpg', '.jpeg'))]
-    if not files:
-        return None
-    latest_file = max(files, key=lambda f: os.path.getmtime(os.path.join(image_folder, f)))
-    return latest_file
-
-
 @app.route('/')
 def index():
-    latest_image = get_latest_image()
-    return render_template("index.html", image=latest_image)
+    # latest_image = get_latest_image()
+    return render_template("index.html", image="./image.png")
 
 @app.route('/images/<filename>')
 def images(filename):
@@ -32,4 +23,5 @@ def create():
 
 if __name__ == '__main__':
     os.makedirs(image_folder, exist_ok=True)  # Ensure the folder exists
-    app.run(debug=True)
+    app.run(host="0.0.0.0", port=5000, debug=True)
+
diff --git a/lib.py b/lib.py
index 75653dc..86bc982 100644
--- a/lib.py
+++ b/lib.py
@@ -4,8 +4,7 @@ import logging
 import sys
 import litellm
 import time
-
-from datetime import datetime
+import os
 
 from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper
 
@@ -19,6 +18,20 @@ except KeyError as 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"],
@@ -55,6 +68,7 @@ def generate_image(file_name, comfy_prompt):
         # 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+"
@@ -69,8 +83,4 @@ def create_image():
     """Main function for generating images."""
     prompt = send_prompt_to_openwebui(user_config["comfyui"]["prompt"])
     print(f"Generated prompt: {prompt}")
-    generate_image(str(time.time()), prompt)
-
-
-# if __name__ == "__main__":
-#     main()
+    generate_image("image", prompt)