From e4428b46928ed5d80c323af9da955d6f9f6fad71 Mon Sep 17 00:00:00 2001
From: Karl Hudgell <karl.hudgell@bjss.com>
Date: Sat, 12 Apr 2025 09:50:15 +0100
Subject: [PATCH] return generated prompt, allow job to be cancelled

---
 ai_frame_image_server.py | 18 +++++++++++++++---
 lib.py                   | 14 ++++++++++++--
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/ai_frame_image_server.py b/ai_frame_image_server.py
index 93fe687..732de54 100644
--- a/ai_frame_image_server.py
+++ b/ai_frame_image_server.py
@@ -2,8 +2,6 @@ from flask import (
     Flask,
     render_template,
     send_from_directory,
-    redirect,
-    url_for,
     request,
     jsonify,
 )
@@ -11,7 +9,7 @@ import os
 import time
 import threading
 from apscheduler.schedulers.background import BackgroundScheduler
-from lib import create_image, load_config
+from lib import create_image, load_config, create_prompt_on_openwebui, cancel_current_job
 
 user_config = load_config()
 app = Flask(__name__)
@@ -59,6 +57,17 @@ def images(filename: str) -> None:
     return send_from_directory(image_folder, filename)
 
 
+@app.route("/cancel", methods=["GET"])
+def cancel_job() -> None:
+    """
+    Serves the requested image file.
+    Args:
+        filename (str): The name of the image file.
+    Returns:
+        None: Sends the image file.
+    """
+    return cancel_current_job()
+
 
 @app.route("/create", methods=["GET", "POST"])
 def create() -> str:
@@ -70,6 +79,9 @@ def create() -> str:
     """
     prompt = request.form.get("prompt") if request.method == "POST" else None
 
+    if prompt is None:
+        prompt = create_prompt_on_openwebui(user_config["comfyui"]["prompt"])
+        
     def create_image_in_background():
         create_image(prompt)
 
diff --git a/lib.py b/lib.py
index 1d72f87..d532226 100644
--- a/lib.py
+++ b/lib.py
@@ -31,6 +31,16 @@ def get_available_models() -> list:
         return []
 
 
+def cancel_current_job() -> list:
+    """Fetches available models from ComfyUI."""
+    url = user_config["comfyui"]["comfyui_url"] + "/interrupt"
+    response = requests.post(url)
+    if response.status_code == 200:
+        return "Cancelled"
+    else:
+        return "Failed to cancel"
+
+
 def load_config() -> configparser.ConfigParser:
     """Loads user configuration from ./user_config.cfg."""
     user_config = configparser.ConfigParser()
@@ -58,7 +68,7 @@ def rename_image() -> str | None:
         return None
 
 
-def send_prompt_to_openwebui(prompt: str) -> str:
+def create_prompt_on_openwebui(prompt: str) -> str:
     """Sends prompt to OpenWebui and returns the generated response."""
     response = litellm.completion(
         api_base=user_config["openwebui"]["base_url"],
@@ -164,7 +174,7 @@ def generate_image(file_name: str, comfy_prompt: str) -> None:
 def create_image(prompt: str | None = None) -> None:
     """Main function for generating images."""
     if prompt is None:
-        prompt = send_prompt_to_openwebui(user_config["comfyui"]["prompt"])
+        prompt = create_prompt_on_openwebui(user_config["comfyui"]["prompt"])
     if prompt:
         logging.info(f"Generated prompt: {prompt}")  # Log generated prompt
         generate_image("image", prompt)