mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-17 08:52:34 +08:00
273 lines
10 KiB
Python
273 lines
10 KiB
Python
import uuid
|
|
import asyncio
|
|
import contextlib
|
|
from io import BytesIO
|
|
import logging
|
|
import time
|
|
from typing import Optional, Union
|
|
|
|
import aiohttp
|
|
import torch
|
|
from pydantic import BaseModel, Field
|
|
|
|
from comfy_api.latest import IO
|
|
from urllib.parse import urlparse
|
|
from .api_client import (
|
|
ApiEndpoint,
|
|
sync_op_pydantic,
|
|
_display_time_progress,
|
|
_diagnose_connectivity,
|
|
)
|
|
|
|
from comfy_api_nodes.apis import request_logger
|
|
from comfy_api_nodes.apinode_utils import tensor_to_bytesio
|
|
from ._helpers import _sleep_with_interrupt, _is_processing_interrupted
|
|
from .common_exceptions import ProcessingInterrupted, LocalNetworkError, ApiServerError
|
|
|
|
|
|
class UploadRequest(BaseModel):
|
|
file_name: str = Field(..., description="Filename to upload")
|
|
content_type: Optional[str] = Field(
|
|
None,
|
|
description="Mime type of the file. For example: image/png, image/jpeg, video/mp4, etc.",
|
|
)
|
|
|
|
|
|
class UploadResponse(BaseModel):
|
|
download_url: str = Field(..., description="URL to GET uploaded file")
|
|
upload_url: str = Field(..., description="URL to PUT file to upload")
|
|
|
|
|
|
async def upload_images_to_comfyapi(
|
|
cls: type[IO.ComfyNode],
|
|
image: torch.Tensor,
|
|
*,
|
|
max_images: int = 8,
|
|
mime_type: Optional[str] = None,
|
|
wait_label: Optional[str] = "Uploading",
|
|
) -> list[str]:
|
|
"""
|
|
Uploads images to ComfyUI API and returns download URLs.
|
|
To upload multiple images, stack them in the batch dimension first.
|
|
"""
|
|
# if batch, try to upload each file if max_images is greater than 0
|
|
download_urls: list[str] = []
|
|
is_batch = len(image.shape) > 3
|
|
batch_len = image.shape[0] if is_batch else 1
|
|
|
|
for idx in range(min(batch_len, max_images)):
|
|
tensor = image[idx] if is_batch else image
|
|
img_io = tensor_to_bytesio(tensor, mime_type=mime_type)
|
|
url = await upload_file_to_comfyapi(cls, img_io, img_io.name, mime_type, wait_label)
|
|
download_urls.append(url)
|
|
return download_urls
|
|
|
|
|
|
async def upload_file_to_comfyapi(
|
|
cls: type[IO.ComfyNode],
|
|
file_bytes_io: BytesIO,
|
|
filename: str,
|
|
upload_mime_type: Optional[str],
|
|
wait_label: Optional[str] = "Uploading",
|
|
) -> str:
|
|
"""Uploads a single file to ComfyUI API and returns its download URL."""
|
|
if upload_mime_type is None:
|
|
request_object = UploadRequest(file_name=filename)
|
|
else:
|
|
request_object = UploadRequest(file_name=filename, content_type=upload_mime_type)
|
|
create_resp = await sync_op_pydantic(
|
|
cls,
|
|
endpoint=ApiEndpoint(path="/customers/storage", method="POST"),
|
|
data=request_object,
|
|
response_model=UploadResponse,
|
|
final_label_on_success=None,
|
|
monitor_progress=False,
|
|
)
|
|
await upload_file(
|
|
cls, create_resp.upload_url,
|
|
file_bytes_io,
|
|
content_type=upload_mime_type,
|
|
wait_label=wait_label,
|
|
)
|
|
return create_resp.download_url
|
|
|
|
|
|
async def upload_file(
|
|
cls: type[IO.ComfyNode],
|
|
upload_url: str,
|
|
file: Union[BytesIO, str],
|
|
*,
|
|
content_type: Optional[str] = None,
|
|
max_retries: int = 3,
|
|
retry_delay: float = 1.0,
|
|
retry_backoff: float = 2.0,
|
|
wait_label: Optional[str] = None,
|
|
) -> None:
|
|
"""
|
|
Upload a file to a signed URL (e.g., S3 pre-signed PUT) with retries, Comfy progress display, and interruption.
|
|
|
|
Args:
|
|
cls: Node class (provides auth context + UI progress hooks).
|
|
upload_url: Pre-signed PUT URL.
|
|
file: BytesIO or path string.
|
|
content_type: Explicit MIME type. If None, we *suppress* Content-Type.
|
|
max_retries: Maximum retry attempts.
|
|
retry_delay: Initial delay in seconds.
|
|
retry_backoff: Exponential backoff factor.
|
|
wait_label: Progress label shown in Comfy UI.
|
|
|
|
Raises:
|
|
ProcessingInterrupted, LocalNetworkError, ApiServerError, Exception
|
|
"""
|
|
if isinstance(file, BytesIO):
|
|
with contextlib.suppress(Exception):
|
|
file.seek(0)
|
|
data = file.read()
|
|
elif isinstance(file, str):
|
|
with open(file, "rb") as f:
|
|
data = f.read()
|
|
else:
|
|
raise ValueError("file must be a BytesIO or a filesystem path string")
|
|
|
|
headers: dict[str, str] = {}
|
|
skip_auto_headers: set[str] = set()
|
|
if content_type:
|
|
headers["Content-Type"] = content_type
|
|
else:
|
|
skip_auto_headers.add("Content-Type") # Don't let aiohttp add Content-Type, it can break the signed request
|
|
|
|
attempt = 0
|
|
delay = retry_delay
|
|
start_ts = time.monotonic()
|
|
op_uuid = uuid.uuid4().hex[:8]
|
|
while True:
|
|
attempt += 1
|
|
operation_id = _generate_operation_id("PUT", upload_url, attempt, op_uuid)
|
|
timeout = aiohttp.ClientTimeout(total=None)
|
|
stop_evt = asyncio.Event()
|
|
|
|
async def _monitor():
|
|
try:
|
|
while not stop_evt.is_set():
|
|
if _is_processing_interrupted():
|
|
return
|
|
if wait_label:
|
|
_display_time_progress(cls, wait_label, int(time.monotonic() - start_ts), None)
|
|
await asyncio.sleep(1.0)
|
|
except asyncio.CancelledError:
|
|
return
|
|
|
|
monitor_task = asyncio.create_task(_monitor())
|
|
sess: Optional[aiohttp.ClientSession] = None
|
|
try:
|
|
try:
|
|
request_logger.log_request_response(
|
|
operation_id=operation_id,
|
|
request_method="PUT",
|
|
request_url=upload_url,
|
|
request_headers=headers or None,
|
|
request_params=None,
|
|
request_data=f"[File data {len(data)} bytes]",
|
|
)
|
|
except Exception as e:
|
|
logging.debug("[DEBUG] upload request logging failed: %s", e)
|
|
|
|
sess = aiohttp.ClientSession(timeout=timeout)
|
|
req = sess.put(upload_url, data=data, headers=headers, skip_auto_headers=skip_auto_headers)
|
|
req_task = asyncio.create_task(req)
|
|
|
|
done, pending = await asyncio.wait({req_task, monitor_task}, return_when=asyncio.FIRST_COMPLETED)
|
|
|
|
if monitor_task in done and req_task in pending:
|
|
req_task.cancel()
|
|
raise ProcessingInterrupted("Upload cancelled")
|
|
|
|
resp = await req_task
|
|
async with resp:
|
|
if resp.status >= 400:
|
|
with contextlib.suppress(Exception):
|
|
try:
|
|
body = await resp.json()
|
|
except Exception:
|
|
body = await resp.text()
|
|
msg = f"Upload failed with status {resp.status}"
|
|
request_logger.log_request_response(
|
|
operation_id=operation_id,
|
|
request_method="PUT",
|
|
request_url=upload_url,
|
|
response_status_code=resp.status,
|
|
response_headers=dict(resp.headers),
|
|
response_content=body,
|
|
error_message=msg,
|
|
)
|
|
if resp.status in {408, 429, 500, 502, 503, 504} and attempt <= max_retries:
|
|
await _sleep_with_interrupt(
|
|
delay,
|
|
cls,
|
|
wait_label,
|
|
start_ts,
|
|
None,
|
|
display_callback=_display_time_progress if wait_label else None,
|
|
)
|
|
delay *= retry_backoff
|
|
continue
|
|
raise Exception(f"Failed to upload (HTTP {resp.status}).")
|
|
try:
|
|
request_logger.log_request_response(
|
|
operation_id=operation_id,
|
|
request_method="PUT",
|
|
request_url=upload_url,
|
|
response_status_code=resp.status,
|
|
response_headers=dict(resp.headers),
|
|
response_content="File uploaded successfully.",
|
|
)
|
|
except Exception as e:
|
|
logging.debug("[DEBUG] upload response logging failed: %s", e)
|
|
return
|
|
except (aiohttp.ClientError, asyncio.TimeoutError) as e:
|
|
if attempt <= max_retries:
|
|
with contextlib.suppress(Exception):
|
|
request_logger.log_request_response(
|
|
operation_id=operation_id,
|
|
request_method="PUT",
|
|
request_url=upload_url,
|
|
request_headers=headers or None,
|
|
request_data=f"[File data {len(data)} bytes]",
|
|
error_message=f"{type(e).__name__}: {str(e)} (will retry)",
|
|
)
|
|
await _sleep_with_interrupt(
|
|
delay,
|
|
cls,
|
|
wait_label,
|
|
start_ts,
|
|
None,
|
|
display_callback=_display_time_progress if wait_label else None,
|
|
)
|
|
delay *= retry_backoff
|
|
continue
|
|
|
|
diag = await _diagnose_connectivity()
|
|
if diag.get("is_local_issue"):
|
|
raise LocalNetworkError(
|
|
"Unable to connect to the network. Please check your internet connection and try again."
|
|
) from e
|
|
raise ApiServerError("The API service appears unreachable at this time.") from e
|
|
finally:
|
|
stop_evt.set()
|
|
if monitor_task:
|
|
monitor_task.cancel()
|
|
with contextlib.suppress(Exception):
|
|
await monitor_task
|
|
if sess:
|
|
with contextlib.suppress(Exception):
|
|
await sess.close()
|
|
|
|
|
|
def _generate_operation_id(method: str, url: str, attempt: int, op_uuid: str) -> str:
|
|
try:
|
|
parsed = urlparse(url)
|
|
slug = (parsed.path.rsplit("/", 1)[-1] or parsed.netloc or "upload").strip("/").replace("/", "_")
|
|
except Exception:
|
|
slug = "upload"
|
|
return f"{method}_{slug}_{op_uuid}_try{attempt}"
|