Add support for Comfy API keys (#8041)

* Handle Comfy API key based authorizaton (#167)

Co-authored-by: Jedrzej Kosinski <kosinkadink1@gmail.com>

* Bump frontend version to include API key features (#170)

* bump templates version

---------

Co-authored-by: Jedrzej Kosinski <kosinkadink1@gmail.com>

Author: christian-byrne

comfy_api_nodes/apinode_utils.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import io
 import logging
 from typing import Optional
@@ -314,7 +315,7 @@ def upload_file_to_comfyapi(
     file_bytes_io: BytesIO,
     filename: str,
     upload_mime_type: str,
-    auth_token: Optional[str] = None,
+    auth_kwargs: Optional[dict[str,str]] = None,
 ) -> str:
     """
     Uploads a single file to ComfyUI API and returns its download URL.
@@ -323,7 +324,7 @@ def upload_file_to_comfyapi(
         file_bytes_io: BytesIO object containing the file data.
         filename: The filename of the file.
         upload_mime_type: MIME type of the file.
-        auth_token: Optional authentication token.
+        auth_kwargs: Optional authentication token(s).
 
     Returns:
         The download URL for the uploaded file.
@@ -337,7 +338,7 @@ def upload_file_to_comfyapi(
             response_model=UploadResponse,
         ),
         request=request_object,
-        auth_token=auth_token,
+        auth_kwargs=auth_kwargs,
     )
 
     response: UploadResponse = operation.execute()
@@ -351,7 +352,7 @@ def upload_file_to_comfyapi(
 
 def upload_video_to_comfyapi(
     video: VideoInput,
-    auth_token: Optional[str] = None,
+    auth_kwargs: Optional[dict[str,str]] = None,
     container: VideoContainer = VideoContainer.MP4,
     codec: VideoCodec = VideoCodec.H264,
     max_duration: Optional[int] = None,
@@ -362,7 +363,7 @@ def upload_video_to_comfyapi(
 
     Args:
         video: VideoInput object (Comfy VIDEO type).
-        auth_token: Optional authentication token.
+        auth_kwargs: Optional authentication token(s).
         container: The video container format to use (default: MP4).
         codec: The video codec to use (default: H264).
         max_duration: Optional maximum duration of the video in seconds. If the video is longer than this, an error will be raised.
@@ -390,7 +391,7 @@ def upload_video_to_comfyapi(
     video_bytes_io.seek(0)
 
     return upload_file_to_comfyapi(
-        video_bytes_io, filename, upload_mime_type, auth_token
+        video_bytes_io, filename, upload_mime_type, auth_kwargs
     )
 
 
@@ -453,7 +454,7 @@ def audio_ndarray_to_bytesio(
 
 def upload_audio_to_comfyapi(
     audio: AudioInput,
-    auth_token: Optional[str] = None,
+    auth_kwargs: Optional[dict[str,str]] = None,
     container_format: str = "mp4",
     codec_name: str = "aac",
     mime_type: str = "audio/mp4",
@@ -465,7 +466,7 @@ def upload_audio_to_comfyapi(
 
     Args:
         audio: a Comfy `AUDIO` type (contains waveform tensor and sample_rate)
-        auth_token: Optional authentication token.
+        auth_kwargs: Optional authentication token(s).
 
     Returns:
         The download URL for the uploaded audio file.
@@ -477,11 +478,11 @@ def upload_audio_to_comfyapi(
         audio_data_np, sample_rate, container_format, codec_name
     )
 
-    return upload_file_to_comfyapi(audio_bytes_io, filename, mime_type, auth_token)
+    return upload_file_to_comfyapi(audio_bytes_io, filename, mime_type, auth_kwargs)
 
 
 def upload_images_to_comfyapi(
-    image: torch.Tensor, max_images=8, auth_token=None, mime_type: Optional[str] = None
+    image: torch.Tensor, max_images=8, auth_kwargs: Optional[dict[str,str]] = None, mime_type: Optional[str] = None
 ) -> list[str]:
     """
     Uploads images to ComfyUI API and returns download URLs.
@@ -490,7 +491,7 @@ def upload_images_to_comfyapi(
     Args:
         image: Input torch.Tensor image.
         max_images: Maximum number of images to upload.
-        auth_token: Optional authentication token.
+        auth_kwargs: Optional authentication token(s).
         mime_type: Optional MIME type for the image.
     """
     # if batch, try to upload each file if max_images is greater than 0
@@ -521,7 +522,7 @@ def upload_images_to_comfyapi(
                 response_model=UploadResponse,
             ),
             request=request_object,
-            auth_token=auth_token,
+            auth_kwargs=auth_kwargs,
         )
         response = operation.execute()
 

comfy_api_nodes/apis/client.py

@@ -20,7 +20,8 @@
 # 1. Create the API client
 api_client = ApiClient(
     base_url="https://api.example.com",
-    api_key="your_api_key_here",
+    auth_token="your_auth_token_here",
+    comfy_api_key="your_comfy_api_key_here",
     timeout=30.0,
     verify_ssl=True
 )
@@ -146,12 +147,14 @@ class ApiClient:
     def __init__(
         self,
         base_url: str,
-        api_key: Optional[str] = None,
+        auth_token: Optional[str] = None,
+        comfy_api_key: Optional[str] = None,
         timeout: float = 3600.0,
         verify_ssl: bool = True,
     ):
         self.base_url = base_url
-        self.api_key = api_key
+        self.auth_token = auth_token
+        self.comfy_api_key = comfy_api_key
         self.timeout = timeout
         self.verify_ssl = verify_ssl
 
@@ -201,8 +204,10 @@ def get_headers(self) -> Dict[str, str]:
         """Get headers for API requests, including authentication if available"""
         headers = {"Content-Type": "application/json", "Accept": "application/json"}
 
-        if self.api_key:
-            headers["Authorization"] = f"Bearer {self.api_key}"
+        if self.auth_token:
+            headers["Authorization"] = f"Bearer {self.auth_token}"
+        elif self.comfy_api_key:
+            headers["X-API-KEY"] = self.comfy_api_key
 
         return headers
 
@@ -236,7 +241,7 @@ def request(
             requests.RequestException: If the request fails
         """
         url = urljoin(self.base_url, path)
-        self.check_auth_token(self.api_key)
+        self.check_auth(self.auth_token, self.comfy_api_key)
         # Combine default headers with any provided headers
         request_headers = self.get_headers()
         if headers:
@@ -320,11 +325,11 @@ def request(
             return response.json()
         return {}
 
-    def check_auth_token(self, auth_token):
-        """Verify that an auth token is present."""
-        if auth_token is None:
+    def check_auth(self, auth_token, comfy_api_key):
+        """Verify that an auth token is present or comfy_api_key is present"""
+        if auth_token is None and comfy_api_key is None:
             raise Exception("Unauthorized: Please login first to use this node.")
-        return auth_token
+        return auth_token or comfy_api_key
 
     @staticmethod
     def upload_file(
@@ -392,6 +397,8 @@ def __init__(
         files: Optional[Dict[str, Any]] = None,
         api_base: str | None = None,
         auth_token: Optional[str] = None,
+        comfy_api_key: Optional[str] = None,
+        auth_kwargs: Optional[Dict[str,str]] = None,
         timeout: float = 604800.0,
         verify_ssl: bool = True,
         content_type: str = "application/json",
@@ -403,6 +410,10 @@ def __init__(
         self.error = None
         self.api_base: str = api_base or args.comfy_api_base
         self.auth_token = auth_token
+        self.comfy_api_key = comfy_api_key
+        if auth_kwargs is not None:
+            self.auth_token = auth_kwargs.get("auth_token", self.auth_token)
+            self.comfy_api_key = auth_kwargs.get("comfy_api_key", self.comfy_api_key)
         self.timeout = timeout
         self.verify_ssl = verify_ssl
         self.files = files
@@ -415,7 +426,8 @@ def execute(self, client: Optional[ApiClient] = None) -> R:
             if client is None:
                 client = ApiClient(
                     base_url=self.api_base,
-                    api_key=self.auth_token,
+                    auth_token=self.auth_token,
+                    comfy_api_key=self.comfy_api_key,
                     timeout=self.timeout,
                     verify_ssl=self.verify_ssl,
                 )
@@ -502,12 +514,18 @@ def __init__(
         request: Optional[T] = None,
         api_base: str | None = None,
         auth_token: Optional[str] = None,
+        comfy_api_key: Optional[str] = None,
+        auth_kwargs: Optional[Dict[str,str]] = None,
         poll_interval: float = 5.0,
     ):
         self.poll_endpoint = poll_endpoint
         self.request = request
         self.api_base: str = api_base or args.comfy_api_base
         self.auth_token = auth_token
+        self.comfy_api_key = comfy_api_key
+        if auth_kwargs is not None:
+            self.auth_token = auth_kwargs.get("auth_token", self.auth_token)
+            self.comfy_api_key = auth_kwargs.get("comfy_api_key", self.comfy_api_key)
         self.poll_interval = poll_interval
 
         # Polling configuration
@@ -528,7 +546,8 @@ def execute(self, client: Optional[ApiClient] = None) -> R:
             if client is None:
                 client = ApiClient(
                     base_url=self.api_base,
-                    api_key=self.auth_token,
+                    auth_token=self.auth_token,
+                    comfy_api_key=self.comfy_api_key,
                 )
             return self._poll_until_complete(client)
         except Exception as e:

comfy_api_nodes/nodes_bfl.py

@@ -179,6 +179,7 @@ def INPUT_TYPES(s):
             },
             "hidden": {
                 "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
             },
         }
 
@@ -211,7 +212,6 @@ def api_call(
         seed=0,
         image_prompt=None,
         image_prompt_strength=0.1,
-        auth_token=None,
         **kwargs,
     ):
         if image_prompt is None:
@@ -244,7 +244,7 @@ def api_call(
                     None if image_prompt is None else round(image_prompt_strength, 2)
                 ),
             ),
-            auth_token=auth_token,
+            auth_kwargs=kwargs,
         )
         output_image = handle_bfl_synchronous_operation(operation)
         return (output_image,)
@@ -319,6 +319,7 @@ def INPUT_TYPES(s):
             },
             "hidden": {
                 "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
             },
         }
 
@@ -337,7 +338,6 @@ def api_call(
         seed=0,
         image_prompt=None,
         # image_prompt_strength=0.1,
-        auth_token=None,
         **kwargs,
     ):
         image_prompt = (
@@ -361,7 +361,7 @@ def api_call(
                 seed=seed,
                 image_prompt=image_prompt,
             ),
-            auth_token=auth_token,
+            auth_kwargs=kwargs,
         )
         output_image = handle_bfl_synchronous_operation(operation)
         return (output_image,)
@@ -461,6 +461,7 @@ def INPUT_TYPES(s):
             },
             "hidden": {
                 "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
             },
         }
 
@@ -482,7 +483,6 @@ def api_call(
         steps: int,
         guidance: float,
         seed=0,
-        auth_token=None,
         **kwargs,
     ):
         image = convert_image_to_base64(image)
@@ -506,7 +506,7 @@ def api_call(
                 seed=seed,
                 image=image,
             ),
-            auth_token=auth_token,
+            auth_kwargs=kwargs,
         )
         output_image = handle_bfl_synchronous_operation(operation)
         return (output_image,)
@@ -572,6 +572,7 @@ def INPUT_TYPES(s):
             },
             "hidden": {
                 "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
             },
         }
 
@@ -590,7 +591,6 @@ def api_call(
         steps: int,
         guidance: float,
         seed=0,
-        auth_token=None,
         **kwargs,
     ):
         # prepare mask
@@ -615,7 +615,7 @@ def api_call(
                 image=image,
                 mask=mask,
             ),
-            auth_token=auth_token,
+            auth_kwargs=kwargs,
         )
         output_image = handle_bfl_synchronous_operation(operation)
         return (output_image,)
@@ -706,6 +706,7 @@ def INPUT_TYPES(s):
             },
             "hidden": {
                 "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
             },
         }
 
@@ -726,7 +727,6 @@ def api_call(
         steps: int,
         guidance: float,
         seed=0,
-        auth_token=None,
         **kwargs,
     ):
         control_image = convert_image_to_base64(control_image[:,:,:,:3])
@@ -763,7 +763,7 @@ def scale_value(value: float, min_val=0, max_val=500):
                 canny_high_threshold=canny_high_threshold,
                 preprocessed_image=preprocessed_image,
             ),
-            auth_token=auth_token,
+            auth_kwargs=kwargs,
         )
         output_image = handle_bfl_synchronous_operation(operation)
         return (output_image,)
@@ -834,6 +834,7 @@ def INPUT_TYPES(s):
             },
             "hidden": {
                 "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
             },
         }
 
@@ -852,7 +853,6 @@ def api_call(
         steps: int,
         guidance: float,
         seed=0,
-        auth_token=None,
         **kwargs,
     ):
         control_image = convert_image_to_base64(control_image[:,:,:,:3])
@@ -878,7 +878,7 @@ def api_call(
                 control_image=control_image,
                 preprocessed_image=preprocessed_image,
             ),
-            auth_token=auth_token,
+            auth_kwargs=kwargs,
         )
         output_image = handle_bfl_synchronous_operation(operation)
         return (output_image,)