44import time
55import uuid
66from io import BytesIO
7- from typing import Optional , Union
7+ from typing import Optional
88from urllib .parse import urlparse
99
1010import aiohttp
@@ -48,8 +48,9 @@ async def upload_images_to_comfyapi(
4848 image : torch .Tensor ,
4949 * ,
5050 max_images : int = 8 ,
51- mime_type : Optional [str ] = None ,
52- wait_label : Optional [str ] = "Uploading" ,
51+ mime_type : str | None = None ,
52+ wait_label : str | None = "Uploading" ,
53+ show_batch_index : bool = True ,
5354) -> list [str ]:
5455 """
5556 Uploads images to ComfyUI API and returns download URLs.
@@ -59,11 +60,18 @@ async def upload_images_to_comfyapi(
5960 download_urls : list [str ] = []
6061 is_batch = len (image .shape ) > 3
6162 batch_len = image .shape [0 ] if is_batch else 1
63+ num_to_upload = min (batch_len , max_images )
64+ batch_start_ts = time .monotonic ()
6265
63- for idx in range (min ( batch_len , max_images ) ):
66+ for idx in range (num_to_upload ):
6467 tensor = image [idx ] if is_batch else image
6568 img_io = tensor_to_bytesio (tensor , mime_type = mime_type )
66- url = await upload_file_to_comfyapi (cls , img_io , img_io .name , mime_type , wait_label )
69+
70+ effective_label = wait_label
71+ if wait_label and show_batch_index and num_to_upload > 1 :
72+ effective_label = f"{ wait_label } ({ idx + 1 } /{ num_to_upload } )"
73+
74+ url = await upload_file_to_comfyapi (cls , img_io , img_io .name , mime_type , effective_label , batch_start_ts )
6775 download_urls .append (url )
6876 return download_urls
6977
@@ -126,8 +134,9 @@ async def upload_file_to_comfyapi(
126134 cls : type [IO .ComfyNode ],
127135 file_bytes_io : BytesIO ,
128136 filename : str ,
129- upload_mime_type : Optional [str ],
130- wait_label : Optional [str ] = "Uploading" ,
137+ upload_mime_type : str | None ,
138+ wait_label : str | None = "Uploading" ,
139+ progress_origin_ts : float | None = None ,
131140) -> str :
132141 """Uploads a single file to ComfyUI API and returns its download URL."""
133142 if upload_mime_type is None :
@@ -148,34 +157,26 @@ async def upload_file_to_comfyapi(
148157 file_bytes_io ,
149158 content_type = upload_mime_type ,
150159 wait_label = wait_label ,
160+ progress_origin_ts = progress_origin_ts ,
151161 )
152162 return create_resp .download_url
153163
154164
155165async def upload_file (
156166 cls : type [IO .ComfyNode ],
157167 upload_url : str ,
158- file : Union [ BytesIO , str ] ,
168+ file : BytesIO | str ,
159169 * ,
160- content_type : Optional [ str ] = None ,
170+ content_type : str | None = None ,
161171 max_retries : int = 3 ,
162172 retry_delay : float = 1.0 ,
163173 retry_backoff : float = 2.0 ,
164- wait_label : Optional [str ] = None ,
174+ wait_label : str | None = None ,
175+ progress_origin_ts : float | None = None ,
165176) -> None :
166177 """
167178 Upload a file to a signed URL (e.g., S3 pre-signed PUT) with retries, Comfy progress display, and interruption.
168179
169- Args:
170- cls: Node class (provides auth context + UI progress hooks).
171- upload_url: Pre-signed PUT URL.
172- file: BytesIO or path string.
173- content_type: Explicit MIME type. If None, we *suppress* Content-Type.
174- max_retries: Maximum retry attempts.
175- retry_delay: Initial delay in seconds.
176- retry_backoff: Exponential backoff factor.
177- wait_label: Progress label shown in Comfy UI.
178-
179180 Raises:
180181 ProcessingInterrupted, LocalNetworkError, ApiServerError, Exception
181182 """
@@ -198,7 +199,7 @@ async def upload_file(
198199
199200 attempt = 0
200201 delay = retry_delay
201- start_ts = time .monotonic ()
202+ start_ts = progress_origin_ts if progress_origin_ts is not None else time .monotonic ()
202203 op_uuid = uuid .uuid4 ().hex [:8 ]
203204 while True :
204205 attempt += 1
0 commit comments