add support for multiple models and posting a prompt to /create

ai_frame_image_server.py

@@ -1,4 +1,4 @@
-from flask import Flask, render_template, send_from_directory, redirect, url_for
+from flask import Flask, render_template, send_from_directory, redirect, url_for, request, jsonify
 import os
 from lib import create_image, load_config
 
@@ -10,19 +10,23 @@ image_folder = "./output"
 
 @app.route('/')
 def index():
-    # latest_image = get_latest_image()
     return render_template("index.html", image="./image.png", reload_interval=user_config["frame"]["reload_interval"])
 
 @app.route('/images/<filename>')
 def images(filename):
     return send_from_directory(image_folder, filename)
 
-@app.route('/create')
+@app.route('/create', methods=["GET", "POST"])
 def create():
-    """Endpoint to create a new image."""
-    create_image()
+    """Endpoint to create a new image. Supports optional prompt via POST."""
+    prompt = request.form.get("prompt") if request.method == "POST" else None
+    create_image(prompt)  # Pass prompt to create_image()
+    
+    if request.method == "POST":
+        return jsonify({"message": "Image created", "prompt": prompt}), 200
     return redirect(url_for("index"))
 
+
 if __name__ == '__main__':
     os.makedirs(image_folder, exist_ok=True)  # Ensure the folder exists
     app.run(host="0.0.0.0", port=user_config["frame"]["port"], debug=True)

lib.py

@@ -5,9 +5,19 @@ import sys
 import litellm
 import time
 import os
-
+import requests
 from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper
 
+def get_available_models():
+    url = user_config["comfyui"]["comfyui_url"] + "/object_info"
+    response = requests.get(url)
+    if response.status_code == 200:
+        data = response.json()
+        return data.get("CheckpointLoaderSimple", {}).get("input", {}).get("required", {}).get("ckpt_name", [])[0]
+    else:
+        print(f"Failed to fetch models: {response.status_code}")
+        return []
+
 def load_config():
     user_config = configparser.ConfigParser()
     try:
@@ -18,6 +28,7 @@ def load_config():
         logging.error(f"Missing configuration key: {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")
@@ -32,6 +43,7 @@ def rename_image():
         print("No image.png found.")
         return None
 
+
 def send_prompt_to_openwebui(prompt):
     response = litellm.completion(
         api_base=user_config["openwebui"]["base_url"],
@@ -62,8 +74,12 @@ def generate_image(file_name, comfy_prompt):
         wf.set_node_param("Save Image", "filename_prefix", file_name)
         wf.set_node_param("Empty Latent Image", "width", user_config["comfyui"]["width"])
         wf.set_node_param("Empty Latent Image", "height", user_config["comfyui"]["height"])
+        valid_models = list(set(get_available_models()) & set(user_config["comfyui"]["models"].split(",")))
+        if not valid_models:
+            raise Exception("No valid options available.")
+        model = random.choice(valid_models)
         wf.set_node_param(
-            "Load Checkpoint", "ckpt_name", user_config["comfyui"]["model"]
+            "Load Checkpoint", "ckpt_name", model
         )
         # Queue your workflow for completion
         logging.debug(f"Generating image: {file_name}")
@@ -79,8 +95,9 @@ def generate_image(file_name, comfy_prompt):
         logging.error(f"Failed to generate image for UID: {file_name}. Error: {e}")
 
 
-def create_image():
+def create_image(prompt):
     """Main function for generating images."""
+    if prompt is None:
         prompt = send_prompt_to_openwebui(user_config["comfyui"]["prompt"])
     print(f"Generated prompt: {prompt}")
     generate_image("image", prompt)

user_config.cfg.sample

@@ -4,7 +4,7 @@ port = 5000
 
 [comfyui]
 comfyui_url = http://comfyui
-model = zavychromaxl_v100.safetensors
+models = zavychromaxl_v100.safetensors
 output_dir = ./output/
 prompt = "Be explicit, only return the prompt and no other text, Generate a random detailed prompt for stable diffusion."
 width = 1568