wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-11 14:50:49 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

Fix CLI args issues

Browse Source
This commit is contained in:
doctorpangloss 2024-02-09 01:20:57 -08:00
parent bdc843ced1
commit a3f9d007d4
6 changed files with 27 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
comfy/cli_args.py
View File

@ -122,11 +122,18 @@ parser.add_argument("--analytics-use-identity-provider", action="store_true",
parser.add_argument("--distributed-queue-connection-uri", type=str, default=None,
help="EXAMPLE: \"amqp://guest:guest@127.0.0.1\" - Servers and clients will connect to this AMPQ URL to form a distributed queue and exchange prompt execution requests and progress updates.")
parser.add_argument(
'--distributed-queue-roles',
action='append',
choices=['worker', 'frontend'],
default=["worker", "frontend"],
help='Specifies one or more roles for the distributed queue. Acceptable values are "worker" or "frontend", or both by writing the flag twice with each role. Prompters will start the web UI and connect to the provided AMQP URL to submit prompts; workers will pull requests off the AMQP URL.'
'--distributed-queue-worker',
required=False,
type=bool,
default=True,
help='Workers will pull requests off the AMQP URL.'
)
parser.add_argument(
'--distributed-queue-frontend',
required=False,
type=bool,
default=True,
help='Frontends will start the web UI and connect to the provided AMQP URL to submit prompts.'
)
parser.add_argument("--distributed-queue-name", type=str, default="comfyui",
help="This name will be used by the frontends and workers to exchange prompt requests and replies. Progress updates will be prefixed by the queue name, followed by a '.', then the user ID")

6
comfy/cli_args_types.py
View File

@ -71,7 +71,8 @@ class Configuration(dict):
write_out_config_file (bool): Enable writing out the configuration file.
create_directories (bool): Creates the default models/, input/, output/ and temp/ directories, then exits.
distributed_queue_connection_uri (Optional[str]): Servers and clients will connect to this AMQP URL to form a distributed queue and exchange prompt execution requests and progress updates.
distributed_queue_roles (List[str]): Specifies one or more roles for the distributed queue. Acceptable values are "worker" or "frontend", or both by writing the flag twice with each role. Frontends will start the web UI and connect to the provided AMQP URL to submit prompts; workers will pull requests off the AMQP URL.
distributed_queue_frontend (bool): Frontends will start the web UI and connect to the provided AMQP URL to submit prompts.
distributed_queue_worker (bool): Workers will pull requests off the AMQP URL.
distributed_queue_name (str): This name will be used by the frontends and workers to exchange prompt requests and replies. Progress updates will be prefixed by the queue name, followed by a '.', then the user ID.
"""
def __init__(self, **kwargs):
@ -131,7 +132,8 @@ class Configuration(dict):
self.write_out_config_file: bool = False
self.create_directories: bool = False
self.distributed_queue_connection_uri: Optional[str] = None
self.distributed_queue_roles: List[str] = ["worker", "frontend"]
self.distributed_queue_worker: bool = True
self.distributed_queue_frontend: bool = True
self.distributed_queue_name: str = "comfyui"
for key, value in kwargs.items():
self[key] = value

8
comfy/cmd/main.py
View File

@ -227,10 +227,10 @@ async def main():
if args.distributed_queue_connection_uri is not None:
distributed = True
q = DistributedPromptQueue(
caller_server=server if "frontend" in args.distributed_queue_roles else None,
caller_server=server if args.distributed_queue_frontend else None,
connection_uri=args.distributed_queue_connection_uri,
is_caller="frontend" in args.distributed_queue_roles,
is_callee="worker" in args.distributed_queue_roles,
is_caller=args.distributed_queue_frontend,
is_callee=args.distributed_queue_worker,
loop=loop,
queue_name=args.distributed_queue_name
)
@ -257,7 +257,7 @@ async def main():
# in a distributed setting, the default prompt worker will not be able to send execution events via the websocket
worker_thread_server = server if not distributed else ServerStub()
if not distributed or "worker" in args.distributed_queue_roles:
if not distributed or args.distributed_queue_worker:
if distributed:
logging.warning(f"Distributed workers started in the default thread loop cannot notify clients of progress updates. Instead of comfyui or main.py, use comfyui-worker.")
threading.Thread(target=prompt_worker, daemon=True, args=(q, worker_thread_server,)).start()

1
comfy/cmd/server.py
View File

@ -108,6 +108,7 @@ class PromptServer(ExecutorToClientProgress):
routes = web.RouteTableDef()
self.routes: web.RouteTableDef = routes
self.last_node_id = None
self.last_prompt_id = None
self.client_id = None
self.on_prompt_handlers = []

3
comfy/cmd/worker.py
View File

@ -10,7 +10,8 @@ from ..cli_args import args
async def main():
# assume we are a worker
args.distributed_queue_roles = ["worker"]
args.distributed_queue_worker = True
args.distributed_queue_frontend = False
assert args.distributed_queue_connection_uri is not None, "Set the --distributed-queue-connection-uri argument to your RabbitMQ server"
async with DistributedPromptWorker(connection_uri=args.distributed_queue_connection_uri,

5
comfy/distributed/distributed_prompt_queue.py
View File

@ -131,7 +131,10 @@ class DistributedPromptQueue(AbstractPromptQueue):
assert self._is_callee
# the loop receiving messages must not be mounted on the worker thread
# otherwise receiving messages will be blocked forever
worker_event_loop = asyncio.get_event_loop()
try:
worker_event_loop = asyncio.get_event_loop()
except RuntimeError:
worker_event_loop = None
assert self._loop != worker_event_loop, "get only makes sense in the context of the legacy comfyui prompt worker"
# spin wait
timeout = timeout or 30.0