Fix image upload for Luma: only include Content-Type header field if it's set explicitly (#40)

This commit is contained in:
Christian Byrne 2025-04-29 18:31:09 -07:00 committed by Robin Huang
parent 22c7bacffc
commit 65eb4104a9

View File

@ -226,6 +226,9 @@ def _tensor_to_pil(image: torch.Tensor, total_pixels: int = 2048 * 2048) -> Imag
def _pil_to_bytesio(img: Image.Image, mime_type: str = "image/png") -> BytesIO:
"""Converts a PIL Image to a BytesIO object."""
if not mime_type:
mime_type = "image/png"
img_byte_arr = io.BytesIO()
# Derive PIL format from MIME type (e.g., 'image/png' -> 'PNG')
pil_format = mime_type.split("/")[-1].upper()
@ -253,6 +256,9 @@ def tensor_to_bytesio(
Returns:
Named BytesIO object containing the image data.
"""
if not mime_type:
mime_type = "image/png"
pil_image = _tensor_to_pil(image, total_pixels=total_pixels)
img_binary = _pil_to_bytesio(pil_image, mime_type=mime_type)
img_binary.name = (
@ -304,7 +310,7 @@ def tensor_to_data_uri(
def upload_images_to_comfyapi(
image: torch.Tensor, max_images=8, auth_token=None, mime_type: str = "image/png"
image: torch.Tensor, max_images=8, auth_token=None, mime_type: Optional[str] = None
) -> list[str]:
# if batch, try to upload each file if max_images is greater than 0
idx_image = 0
@ -320,6 +326,12 @@ def upload_images_to_comfyapi(
# get BytesIO version of image
img_binary = tensor_to_bytesio(curr_image, mime_type=mime_type)
# first, request upload/download urls from comfy API
if not mime_type:
request_object = UploadRequest(filename=img_binary.name)
else:
request_object = UploadRequest(
filename=img_binary.name, content_type=mime_type
)
operation = SynchronousOperation(
endpoint=ApiEndpoint(
path="/customers/storage",
@ -327,7 +339,7 @@ def upload_images_to_comfyapi(
request_model=UploadRequest,
response_model=UploadResponse,
),
request=UploadRequest(filename=img_binary.name, content_type=mime_type),
request=request_object,
auth_token=auth_token,
)
response = operation.execute()