diff --git a/comfy_api_nodes/nodes_bytedance.py b/comfy_api_nodes/nodes_bytedance.py index d97e0a3a8..f2e3e9027 100644 --- a/comfy_api_nodes/nodes_bytedance.py +++ b/comfy_api_nodes/nodes_bytedance.py @@ -409,7 +409,7 @@ class ByteDanceImageEditNode(IO.ComfyNode): if get_number_of_images(image) != 1: raise ValueError("Exactly one input image is required.") validate_image_aspect_ratio_range(image, (1, 3), (3, 1)) - source_url = (await upload_images_to_comfyapi(cls, image, max_images=1, mime_type="image/png",))[0] + source_url = (await upload_images_to_comfyapi(cls, image, max_images=1, mime_type="image/png"))[0] payload = Image2ImageTaskCreationRequest( model=model, prompt=prompt, diff --git a/comfy_api_nodes/util/__init__.py b/comfy_api_nodes/util/__init__.py index cad902fc0..fe3cda258 100644 --- a/comfy_api_nodes/util/__init__.py +++ b/comfy_api_nodes/util/__init__.py @@ -1,5 +1,10 @@ from .api_client import ApiEndpoint, sync_op_pydantic, poll_op_pydantic, sync_op, poll_op -from .storage_helpers import ( +from .download_helpers import ( + download_url_to_bytesio, + download_url_to_image_tensor, + bytesio_to_image_tensor, +) +from .upload_helpers import ( upload_file_to_comfyapi, upload_images_to_comfyapi, ) @@ -12,4 +17,7 @@ __all__ = [ "sync_op_pydantic", "upload_file_to_comfyapi", "upload_images_to_comfyapi", + "download_url_to_bytesio", + "download_url_to_image_tensor", + "bytesio_to_image_tensor", ] diff --git a/comfy_api_nodes/util/conversions.py b/comfy_api_nodes/util/conversions.py new file mode 100644 index 000000000..207fe0ef2 --- /dev/null +++ b/comfy_api_nodes/util/conversions.py @@ -0,0 +1,25 @@ +from io import BytesIO + +import numpy as np +from PIL import Image +import torch + + +def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str = "RGBA") -> torch.Tensor: + """Converts image data from BytesIO to a torch.Tensor. + + Args: + image_bytesio: BytesIO object containing the image data. + mode: The PIL mode to convert the image to (e.g., "RGB", "RGBA"). + + Returns: + A torch.Tensor representing the image (1, H, W, C). + + Raises: + PIL.UnidentifiedImageError: If the image data cannot be identified. + ValueError: If the specified mode is invalid. + """ + image = Image.open(image_bytesio) + image = image.convert(mode) + image_array = np.array(image).astype(np.float32) / 255.0 + return torch.from_numpy(image_array).unsqueeze(0) diff --git a/comfy_api_nodes/util/download_helpers.py b/comfy_api_nodes/util/download_helpers.py new file mode 100644 index 000000000..90e127b74 --- /dev/null +++ b/comfy_api_nodes/util/download_helpers.py @@ -0,0 +1,246 @@ +import asyncio +import contextlib +import logging +import time +import uuid +from io import BytesIO +from typing import Optional, Union, IO +from pathlib import Path + +import aiohttp +import torch +from aiohttp.client_exceptions import ClientError, ContentTypeError +from urllib.parse import urlparse + +from comfy_api_nodes.apis import request_logger + +from ._helpers import _is_processing_interrupted +from .common_exceptions import ProcessingInterrupted, LocalNetworkError, ApiServerError +from .api_client import _diagnose_connectivity +from .conversions import bytesio_to_image_tensor + + +_RETRY_STATUS = {408, 429, 500, 502, 503, 504} + + +async def download_url_to_bytesio( + url: str, + timeout: Optional[float] = None, + *, + dest: Optional[Union[BytesIO, IO[bytes], str, Path]] = None, + max_retries: int = 3, + retry_delay: float = 1.0, + retry_backoff: float = 2.0, +) -> None: + """Stream-download a URL into memory or to a provided destination. + + Raises: + ProcessingInterrupted, LocalNetworkError, ApiServerError, Exception (HTTP and other errors) + """ + attempt = 0 + delay = retry_delay + + while True: + attempt += 1 + op_id = _generate_operation_id("GET", url, attempt) + timeout_cfg = aiohttp.ClientTimeout(total=timeout) + stop_evt = asyncio.Event() + + async def _monitor(): + try: + while not stop_evt.is_set(): + if _is_processing_interrupted(): + return + await asyncio.sleep(1.0) + except asyncio.CancelledError: + return + + monitor_task: Optional[asyncio.Task] = None + sess: Optional[aiohttp.ClientSession] = None + + # Open file path if a path was provided + is_path_sink = isinstance(dest, (str, Path)) + fhandle = None + try: + try: + request_logger.log_request_response( + operation_id=op_id, + request_method="GET", + request_url=url, + ) + except Exception as e: + logging.debug("[DEBUG] download request logging failed: %s", e) + + monitor_task = asyncio.create_task(_monitor()) + sess = aiohttp.ClientSession(timeout=timeout_cfg) + req_task = asyncio.create_task(sess.get(url)) + + done, pending = await asyncio.wait({req_task, monitor_task}, return_when=asyncio.FIRST_COMPLETED) + + # Interruption wins the race + if monitor_task in done and req_task in pending: + req_task.cancel() + raise ProcessingInterrupted("Task cancelled") + + resp = await req_task + async with resp: + if resp.status >= 400: + # Attempt to capture body for logging (do not log huge binaries) + with contextlib.suppress(Exception): + try: + body = await resp.json() + except (ContentTypeError, ValueError): + text = await resp.text() + body = text if len(text) <= 4096 else f"[text {len(text)} bytes]" + request_logger.log_request_response( + operation_id=op_id, + request_method="GET", + request_url=url, + response_status_code=resp.status, + response_headers=dict(resp.headers), + response_content=body, + error_message=f"HTTP {resp.status}", + ) + + if resp.status in _RETRY_STATUS and attempt <= max_retries: + await _sleep_with_cancel(delay) + delay *= retry_backoff + continue + raise Exception(f"Failed to download (HTTP {resp.status}).") + + # Prepare path sink if needed + if is_path_sink: + p = Path(str(dest)) + with contextlib.suppress(Exception): + p.parent.mkdir(parents=True, exist_ok=True) + fhandle = open(p, "wb") + sink = fhandle + else: + sink = dest # BytesIO or file-like + + # Stream body in chunks to sink with cancellation checks + written = 0 + last_tick = time.monotonic() + async for chunk in resp.content.iter_chunked(1024 * 1024): + sink.write(chunk) + written += len(chunk) + now = time.monotonic() + if now - last_tick >= 1.0: + last_tick = now + if _is_processing_interrupted(): + raise ProcessingInterrupted("Task cancelled") + + if isinstance(dest, BytesIO): + dest.seek(0) + + try: + request_logger.log_request_response( + operation_id=op_id, + request_method="GET", + request_url=url, + response_status_code=resp.status, + response_headers=dict(resp.headers), + response_content=f"[streamed {written} bytes to dest]", + ) + except Exception as e: + logging.debug("[DEBUG] download response logging failed: %s", e) + return + except ProcessingInterrupted: + logging.debug("Download was interrupted by user") + raise + except (ClientError, asyncio.TimeoutError) as e: + if attempt <= max_retries: + with contextlib.suppress(Exception): + request_logger.log_request_response( + operation_id=op_id, + request_method="GET", + request_url=url, + error_message=f"{type(e).__name__}: {str(e)} (will retry)", + ) + await _sleep_with_cancel(delay) + 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 remote service appears unreachable at this time.") from e + finally: + with contextlib.suppress(Exception): + if fhandle: + fhandle.flush() + fhandle.close() + 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() + + +async def download_url_to_image_tensor( + url: str, + timeout: int = None, + auth_kwargs: Optional[dict[str, str]] = None, + *, + dest: Optional[Union[BytesIO, IO[bytes], str, Path]] = None, + mode: str = "RGBA", +) -> torch.Tensor: + """ + Download image and decode to tensor. Supports streaming `dest` like util version. + """ + if dest is None: + bio = await download_url_to_bytesio(url, timeout, auth_kwargs, dest=None) + return bytesio_to_image_tensor(bio, mode=mode) # type: ignore[arg-type] + + await download_url_to_bytesio(url, timeout, auth_kwargs, dest=dest) + + if isinstance(dest, BytesIO): + with contextlib.suppress(Exception): + dest.seek(0) + return bytesio_to_image_tensor(dest, mode=mode) + + if hasattr(dest, "read") and hasattr(dest, "seek"): + try: + with contextlib.suppress(Exception): + dest.flush() + dest.seek(0) + data = dest.read() + return bytesio_to_image_tensor(BytesIO(data), mode=mode) + except Exception: + pass + + if isinstance(dest, (str, Path)) or getattr(dest, "name", None): + path_str = str(dest if isinstance(dest, (str, Path)) else getattr(dest, "name")) + with open(path_str, "rb") as f: + return bytesio_to_image_tensor(BytesIO(f.read()), mode=mode) + + raise ValueError( + "Destination is not readable and no path is available to decode the image. " + "Pass dest=None to decode from memory, or provide a readable handle / path." + ) + + +def _generate_operation_id(method: str, url: str, attempt: int) -> str: + try: + parsed = urlparse(url) + slug = (parsed.path.rsplit("/", 1)[-1] or parsed.netloc or "download").strip("/").replace("/", "_") + except Exception: + slug = "download" + return f"{method}_{slug}_try{attempt}_{uuid.uuid4().hex[:8]}" + + +async def _sleep_with_cancel(seconds: float) -> None: + """Sleep in 1s slices while checking for interruption.""" + end = time.monotonic() + seconds + while True: + if _is_processing_interrupted(): + raise ProcessingInterrupted("Task cancelled") + now = time.monotonic() + if now >= end: + return + await asyncio.sleep(min(1.0, end - now)) diff --git a/comfy_api_nodes/util/upload_helpers.py b/comfy_api_nodes/util/upload_helpers.py new file mode 100644 index 000000000..457a6f11e --- /dev/null +++ b/comfy_api_nodes/util/upload_helpers.py @@ -0,0 +1,272 @@ +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}"