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 GitHub
parent 2d4d2f0dfe
commit c42295c579
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 22 additions and 6 deletions

View File

@ -105,7 +105,7 @@ from typing import (
TypeVar,
Generic,
)
from pydantic import BaseModel, Field, HttpUrl
from pydantic import BaseModel, Field
from enum import Enum
import json
import requests
@ -127,11 +127,15 @@ class EmptyRequest(BaseModel):
class UploadRequest(BaseModel):
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):
download_url: HttpUrl = Field(..., description='URL to GET uploaded file')
upload_url: HttpUrl = Field(..., description='URL to PUT file to upload')
download_url: str = Field(..., description="URL to GET uploaded file")
upload_url: str = Field(..., description="URL to PUT file to upload")
class HttpMethod(str, Enum):

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()