diff --git a/libs/comfyui.py b/libs/comfyui.py index 0d98999..9d61080 100644 --- a/libs/comfyui.py +++ b/libs/comfyui.py @@ -194,3 +194,17 @@ def create_image(prompt: str | None = None, model: str = "Random Image Model") - generate_image("image", comfy_prompt=prompt, model=model) logging.info(f"{selected_workflow} generation started with prompt: {prompt}") + +def get_queue_count() -> int: + """Fetches the current queue count from ComfyUI (pending + running jobs).""" + url = user_config["comfyui"]["comfyui_url"] + "/queue" + try: + response = requests.get(url) + response.raise_for_status() + data = response.json() + pending = len(data.get("queue_pending", [])) + running = len(data.get("queue_running", [])) + return pending + running + except Exception as e: + logging.error(f"Error fetching queue count: {e}") + return 0 diff --git a/routes/create_routes.py b/routes/create_routes.py index 2d77cfa..b54ad01 100644 --- a/routes/create_routes.py +++ b/routes/create_routes.py @@ -1,6 +1,6 @@ from flask import Blueprint, request, render_template, redirect, url_for, session import threading -from libs.comfyui import create_image, select_model, get_available_models +from libs.comfyui import create_image, select_model, get_available_models, get_queue_count from libs.openwebui import create_prompt_on_openwebui from libs.generic import load_models_from_config, load_topics_from_config, load_openrouter_models_from_config, load_openwebui_models_from_config, create_prompt_with_random_model import os @@ -40,13 +40,15 @@ def create(): openwebui_models = load_openwebui_models_from_config() openrouter_models = load_openrouter_models_from_config() + queue_count = get_queue_count() return render_template("create_image.html", - sdxl_models=sdxl_models, + sdxx_models=sdxl_models, flux_models=flux_models, qwen_models=qwen_models, openwebui_models=openwebui_models, openrouter_models=openrouter_models, - topics=load_topics_from_config()) + topics=load_topics_from_config(), + queue_count=queue_count) @bp.route("/image_queued") def image_queued(): @@ -68,13 +70,15 @@ def create_image_page(): openwebui_models = load_openwebui_models_from_config() openrouter_models = load_openrouter_models_from_config() + queue_count = get_queue_count() return render_template("create_image.html", sdxl_models=sdxl_models, flux_models=flux_models, qwen_models=qwen_models, openwebui_models=openwebui_models, openrouter_models=openrouter_models, - topics=load_topics_from_config()) + topics=load_topics_from_config(), + queue_count=queue_count) def init_app(config): diff --git a/templates/create_image.html b/templates/create_image.html index 28e5cea..9c246d4 100644 --- a/templates/create_image.html +++ b/templates/create_image.html @@ -135,6 +135,9 @@ {% endblock %} {% block content %} +