mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-10 14:20:49 +08:00
132 lines
4.9 KiB
Python
132 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import gc
|
|
import uuid
|
|
from asyncio import AbstractEventLoop
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import Optional
|
|
|
|
from ..api.components.schema.prompt import PromptDict
|
|
from ..cli_args_types import Configuration
|
|
from ..component_model.make_mutable import make_mutable
|
|
from ..component_model.executor_types import ExecutorToClientProgress
|
|
from ..distributed.server_stub import ServerStub
|
|
|
|
|
|
class EmbeddedComfyClient:
|
|
"""
|
|
Embedded client for comfy executing prompts as a library.
|
|
|
|
This client manages a single-threaded executor to run long-running or blocking tasks
|
|
asynchronously without blocking the asyncio event loop. It initializes a PromptExecutor
|
|
in a dedicated thread for executing prompts and handling server-stub communications.
|
|
|
|
Example usage:
|
|
|
|
Asynchronous (non-blocking) usage with async-await:
|
|
```
|
|
# Write a workflow, or enable Dev Mode in the UI settings, then Save (API Format) to get the workflow in your
|
|
# workspace.
|
|
prompt_dict = {
|
|
"1": {"class_type": "KSamplerAdvanced", ...}
|
|
...
|
|
}
|
|
|
|
# Validate your workflow (the prompt)
|
|
from comfy.api.components.schema.prompt import Prompt
|
|
prompt = Prompt.validate(prompt_dict)
|
|
# Then use the client to run your workflow. This will start, then stop, a local ComfyUI workflow executor.
|
|
# It does not connect to a remote server.
|
|
async def main():
|
|
async with EmbeddedComfyClient() as client:
|
|
outputs = await client.queue_prompt(prompt)
|
|
print(outputs)
|
|
print("Now that we've exited the with statement, all your VRAM has been cleared from ComfyUI")
|
|
|
|
if __name__ == "__main__"
|
|
asyncio.run(main())
|
|
```
|
|
|
|
In order to use this in blocking methods, learn more about asyncio online.
|
|
"""
|
|
|
|
def __init__(self, configuration: Optional[Configuration] = None,
|
|
progress_handler: Optional[ExecutorToClientProgress] = None,
|
|
loop: Optional[AbstractEventLoop] = None,
|
|
max_workers: int = 1):
|
|
self._progress_handler = progress_handler or ServerStub()
|
|
self._executor = ThreadPoolExecutor(max_workers=max_workers)
|
|
self._loop = loop or asyncio.get_event_loop()
|
|
self._configuration = configuration
|
|
# we don't want to import the executor yet
|
|
self._prompt_executor: Optional["comfy.cmd.execution.PromptExecutor"] = None
|
|
self._is_running = False
|
|
|
|
@property
|
|
def is_running(self) -> bool:
|
|
return self._is_running
|
|
|
|
async def __aenter__(self):
|
|
await self._initialize_prompt_executor()
|
|
self._is_running = True
|
|
return self
|
|
|
|
async def __aexit__(self, *args):
|
|
# Perform cleanup here
|
|
def cleanup():
|
|
from .. import model_management
|
|
model_management.unload_all_models()
|
|
gc.collect()
|
|
try:
|
|
model_management.soft_empty_cache()
|
|
except:
|
|
pass
|
|
|
|
# wait until the queue is done
|
|
while self._executor._work_queue.qsize() > 0:
|
|
await asyncio.sleep(0.1)
|
|
|
|
await self._loop.run_in_executor(self._executor, cleanup)
|
|
|
|
self._executor.shutdown(wait=True)
|
|
self._is_running = False
|
|
|
|
async def _initialize_prompt_executor(self):
|
|
# This method must be async since it's used in __aenter__
|
|
def create_executor_in_thread():
|
|
from .. import options
|
|
if self._configuration is None:
|
|
options.enable_args_parsing()
|
|
else:
|
|
from ..cli_args import args
|
|
args.clear()
|
|
args.update(self._configuration)
|
|
|
|
from ..cmd.execution import PromptExecutor
|
|
|
|
self._prompt_executor = PromptExecutor(self._progress_handler)
|
|
|
|
await self._loop.run_in_executor(self._executor, create_executor_in_thread)
|
|
|
|
async def queue_prompt(self,
|
|
prompt: PromptDict | dict,
|
|
prompt_id: Optional[str] = None,
|
|
client_id: Optional[str] = None) -> dict:
|
|
prompt_id = prompt_id or str(uuid.uuid4())
|
|
client_id = client_id or self._progress_handler.client_id or None
|
|
|
|
def execute_prompt() -> dict:
|
|
from ..cmd.execution import validate_prompt
|
|
prompt_mut = make_mutable(prompt)
|
|
validation_tuple = validate_prompt(prompt_mut)
|
|
|
|
self._prompt_executor.execute(prompt_mut, prompt_id, {"client_id": client_id},
|
|
execute_outputs=validation_tuple[2])
|
|
if self._prompt_executor.success:
|
|
return self._prompt_executor.outputs_ui
|
|
else:
|
|
raise RuntimeError("\n".join(event for (event, data) in self._prompt_executor.status_messages))
|
|
|
|
return await self._loop.run_in_executor(self._executor, execute_prompt)
|