Skip to content

Commit 1105e0d

Browse files
authored
improve UX for batch uploads in upload_images_to_comfyapi (#10913)
1 parent 8938aa3 commit 1105e0d

File tree

1 file changed

+22
-21
lines changed

1 file changed

+22
-21
lines changed

comfy_api_nodes/util/upload_helpers.py

Lines changed: 22 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import time
55
import uuid
66
from io import BytesIO
7-
from typing import Optional, Union
7+
from typing import Optional
88
from urllib.parse import urlparse
99

1010
import 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

155165
async 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

Comments
 (0)