Merge 2700d3bf9c622610e4b433cb4e50f27fed1d0781 into 8af9a91e0c47b9fc277077f2079873adf8edac05

This commit is contained in:
pawelkania 2024-12-06 18:31:34 -05:00 committed by GitHub
commit 25319fa67f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 41 additions and 3 deletions

24
main.py
View File

@ -7,7 +7,8 @@ import folder_paths
import time
from comfy.cli_args import args
from app.logger import setup_logger
import aiohttp
import asyncio
setup_logger(log_level=args.verbose)
@ -105,6 +106,18 @@ def cuda_malloc_warning():
if cuda_malloc_warning:
logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n")
async def send_webhook(server, prompt_id, data):
webhook_url = server.webhooks.pop(prompt_id, None)
if webhook_url:
try:
async with aiohttp.ClientSession() as session:
logging.info(f"Sending webhook for prompt {prompt_id}")
async with session.post(webhook_url, json=data) as response:
if response.status != 200:
logging.warning(f"Webhook delivery failed for prompt {prompt_id}. Status: {response.status}")
except Exception as e:
logging.error(f"Error sending webhook for prompt {prompt_id}: {str(e)}")
def prompt_worker(q, server):
e = execution.PromptExecutor(server, lru_size=args.cache_lru)
last_gc_collect = 0
@ -138,6 +151,15 @@ def prompt_worker(q, server):
execution_time = current_time - execution_start_time
logging.info("Prompt executed in {:.2f} seconds".format(execution_time))
# Send webhook after execution is complete
webhook_data = {
"prompt_id": prompt_id,
"execution_time": execution_time,
"status": "success" if e.success else "error",
"result": e.history_result
}
asyncio.run_coroutine_threadsafe(send_webhook(server, prompt_id, webhook_data), server.loop)
flags = q.get_flags()
free_memory = flags.get("free_memory", False)

View File

@ -14,7 +14,11 @@ scipy
tqdm
psutil
#webhook handling
aiohttp
#non essential dependencies:
kornia>=0.7.1
spandrel
soundfile

View File

@ -31,6 +31,10 @@ from app.frontend_management import FrontendManager
from app.user_manager import UserManager
from typing import Optional
from api_server.routes.internal.internal_routes import InternalRoutes
import requests
import json
from urllib.parse import urlparse, parse_qs, urlencode
class BinaryEventTypes:
PREVIEW_IMAGE = 1
@ -158,6 +162,7 @@ class PromptServer():
self.messages = asyncio.Queue()
self.client_session:Optional[aiohttp.ClientSession] = None
self.number = 0
self.webhooks = {}
middlewares = [cache_control]
if args.enable_cors_header:
@ -622,8 +627,15 @@ class PromptServer():
if "client_id" in json_data:
extra_data["client_id"] = json_data["client_id"]
if valid[0]:
prompt_id = str(uuid.uuid4())
# allow to accept prompt_id from api caller to reference it in webhook handler if needed
prompt_id = json_data.get("prompt_id", str(uuid.uuid4()))
outputs_to_execute = valid[2]
# Add webhook URL to the webhooks dict if provided
webhook_url = json_data.get("webhook_url")
if webhook_url:
self.webhooks[prompt_id] = webhook_url
self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))
response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]}
return web.json_response(response)
@ -833,4 +845,4 @@ class PromptServer():
logging.warning(f"[ERROR] An error occurred during the on_prompt_handler processing")
logging.warning(traceback.format_exc())
return json_data
return json_data