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 Jedrzej Kosinski
parent c997c2c956
commit 1bb11a33fa
2 changed files with 22 additions and 6 deletions

View File

@ -105,7 +105,7 @@ from typing import (
TypeVar, TypeVar,
Generic, Generic,
) )
from pydantic import BaseModel, Field, HttpUrl from pydantic import BaseModel, Field
from enum import Enum from enum import Enum
import json import json
import requests import requests
@ -127,11 +127,15 @@ class EmptyRequest(BaseModel):
class UploadRequest(BaseModel): class UploadRequest(BaseModel):
filename: str = Field(..., description="Filename to upload") filename: str = Field(..., description="Filename to upload")
content_type: str = Field(..., description="Mime type of the file. For example: image/png, image/jpeg, video/mp4, etc.") content_type: str | None = Field(
None,
description="Mime type of the file. For example: image/png, image/jpeg, video/mp4, etc.",
)
class UploadResponse(BaseModel): class UploadResponse(BaseModel):
download_url: HttpUrl = Field(..., description='URL to GET uploaded file') download_url: str = Field(..., description="URL to GET uploaded file")
upload_url: HttpUrl = Field(..., description='URL to PUT file to upload') upload_url: str = Field(..., description="URL to PUT file to upload")
class HttpMethod(str, Enum): class HttpMethod(str, Enum):

View File

@ -229,6 +229,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: def _pil_to_bytesio(img: Image.Image, mime_type: str = "image/png") -> BytesIO:
"""Converts a PIL Image to a BytesIO object.""" """Converts a PIL Image to a BytesIO object."""
if not mime_type:
mime_type = "image/png"
img_byte_arr = io.BytesIO() img_byte_arr = io.BytesIO()
# Derive PIL format from MIME type (e.g., 'image/png' -> 'PNG') # Derive PIL format from MIME type (e.g., 'image/png' -> 'PNG')
pil_format = mime_type.split("/")[-1].upper() pil_format = mime_type.split("/")[-1].upper()
@ -256,6 +259,9 @@ def tensor_to_bytesio(
Returns: Returns:
Named BytesIO object containing the image data. 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) pil_image = _tensor_to_pil(image, total_pixels=total_pixels)
img_binary = _pil_to_bytesio(pil_image, mime_type=mime_type) img_binary = _pil_to_bytesio(pil_image, mime_type=mime_type)
img_binary.name = ( img_binary.name = (
@ -307,7 +313,7 @@ def tensor_to_data_uri(
def upload_images_to_comfyapi( 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]: ) -> list[str]:
# if batch, try to upload each file if max_images is greater than 0 # if batch, try to upload each file if max_images is greater than 0
idx_image = 0 idx_image = 0
@ -323,6 +329,12 @@ def upload_images_to_comfyapi(
# get BytesIO version of image # get BytesIO version of image
img_binary = tensor_to_bytesio(curr_image, mime_type=mime_type) img_binary = tensor_to_bytesio(curr_image, mime_type=mime_type)
# first, request upload/download urls from comfy API # 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( operation = SynchronousOperation(
endpoint=ApiEndpoint( endpoint=ApiEndpoint(
path="/customers/storage", path="/customers/storage",
@ -330,7 +342,7 @@ def upload_images_to_comfyapi(
request_model=UploadRequest, request_model=UploadRequest,
response_model=UploadResponse, response_model=UploadResponse,
), ),
request=UploadRequest(filename=img_binary.name, content_type=mime_type), request=request_object,
auth_token=auth_token, auth_token=auth_token,
) )
response = operation.execute() response = operation.execute()