diff --git a/comfy/cli_args.py b/comfy/cli_args.py index f89a7aab4..a977b02ec 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -192,6 +192,14 @@ parser.add_argument("--user-directory", type=is_valid_directory, default=None, h parser.add_argument("--enable-compress-response-body", action="store_true", help="Enable compressing response body.") +parser.add_argument( + "--comfy-api-base", + type=str, + default="https://api.comfy.org", + choices=["https://api.comfy.org", "https://stagingapi.comfy.org"], + help="Set the base URL for the ComfyUI API.", +) + if comfy.options.args_parsing: args = parser.parse_args() else: diff --git a/comfy_api_nodes/apis/client.py b/comfy_api_nodes/apis/client.py index 49fde8600..dc5768aa9 100644 --- a/comfy_api_nodes/apis/client.py +++ b/comfy_api_nodes/apis/client.py @@ -1,6 +1,9 @@ import logging import time from typing import Callable + +from comfy.cli_args import args + """ API Client Framework for api.comfy.org. @@ -316,7 +319,7 @@ class SynchronousOperation(Generic[T, R]): endpoint: ApiEndpoint[T, R], request: T, files: Optional[Dict[str, Any]] = None, - api_base: str = "https://api.comfy.org", + api_base: str | None = None, auth_token: Optional[str] = None, timeout: float = 604800.0, verify_ssl: bool = True, @@ -325,18 +328,17 @@ class SynchronousOperation(Generic[T, R]): self.request = request self.response = None self.error = None - self.api_base = api_base + self.api_base: str = api_base or args.comfy_api_base self.auth_token = auth_token self.timeout = timeout self.verify_ssl = verify_ssl self.files = files + def execute(self, client: Optional[ApiClient] = None) -> R: """Execute the API operation using the provided client or create one""" try: # Create client if not provided if client is None: - if self.api_base is None: - raise ValueError("Either client or api_base must be provided") client = ApiClient( base_url=self.api_base, api_key=self.auth_token, @@ -410,13 +412,13 @@ class PollingOperation(Generic[T, R]): failed_statuses: list, status_extractor: Callable[[R], str], request: Optional[T] = None, - api_base: str = "https://stagingapi.comfy.org", + api_base: str | None = None, auth_token: Optional[str] = None, poll_interval: float = 1.0, ): self.poll_endpoint = poll_endpoint self.request = request - self.api_base = api_base + self.api_base: str = api_base or args.comfy_api_base self.auth_token = auth_token self.poll_interval = poll_interval