return generated prompt, allow job to be cancelled

This commit is contained in:
Karl Hudgell 2025-04-12 09:50:15 +01:00
parent 91c48b512f
commit e4428b4692
2 changed files with 27 additions and 5 deletions

View File

@ -2,8 +2,6 @@ from flask import (
Flask, Flask,
render_template, render_template,
send_from_directory, send_from_directory,
redirect,
url_for,
request, request,
jsonify, jsonify,
) )
@ -11,7 +9,7 @@ import os
import time import time
import threading import threading
from apscheduler.schedulers.background import BackgroundScheduler 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() user_config = load_config()
app = Flask(__name__) app = Flask(__name__)
@ -59,6 +57,17 @@ def images(filename: str) -> None:
return send_from_directory(image_folder, filename) 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"]) @app.route("/create", methods=["GET", "POST"])
def create() -> str: def create() -> str:
@ -70,6 +79,9 @@ def create() -> str:
""" """
prompt = request.form.get("prompt") if request.method == "POST" else None 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(): def create_image_in_background():
create_image(prompt) create_image(prompt)

14
lib.py
View File

@ -31,6 +31,16 @@ def get_available_models() -> list:
return [] 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: def load_config() -> configparser.ConfigParser:
"""Loads user configuration from ./user_config.cfg.""" """Loads user configuration from ./user_config.cfg."""
user_config = configparser.ConfigParser() user_config = configparser.ConfigParser()
@ -58,7 +68,7 @@ def rename_image() -> str | None:
return 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.""" """Sends prompt to OpenWebui and returns the generated response."""
response = litellm.completion( response = litellm.completion(
api_base=user_config["openwebui"]["base_url"], 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: def create_image(prompt: str | None = None) -> None:
"""Main function for generating images.""" """Main function for generating images."""
if prompt is None: if prompt is None:
prompt = send_prompt_to_openwebui(user_config["comfyui"]["prompt"]) prompt = create_prompt_on_openwebui(user_config["comfyui"]["prompt"])
if prompt: if prompt:
logging.info(f"Generated prompt: {prompt}") # Log generated prompt logging.info(f"Generated prompt: {prompt}") # Log generated prompt
generate_image("image", prompt) generate_image("image", prompt)