fix API type error and add b64 support for 4o

This commit is contained in:
thot-experiment 2025-04-23 11:35:43 -07:00
parent f65b05e865
commit cf209aa456
2 changed files with 16 additions and 8 deletions

View File

@ -573,7 +573,7 @@ class Quality(str, Enum):
class OpenAIImageEditRequest(BaseModel): class OpenAIImageEditRequest(BaseModel):
background: Optional[Background] = Field( background: Optional[str] = Field(
None, description='Background transparency', examples=['opaque'] None, description='Background transparency', examples=['opaque']
) )
model: str = Field( model: str = Field(
@ -596,7 +596,7 @@ class OpenAIImageEditRequest(BaseModel):
description='A text description of the desired edit', description='A text description of the desired edit',
examples=['Give the rocketship rainbow coloring'], examples=['Give the rocketship rainbow coloring'],
) )
quality: Optional[Quality] = Field( quality: Optional[str] = Field(
None, description='The quality of the edited image', examples=['low'] None, description='The quality of the edited image', examples=['low']
) )
size: Optional[str] = Field( size: Optional[str] = Field(

View File

@ -15,6 +15,7 @@ from PIL import Image
import requests import requests
import torch import torch
import math import math
import base64
def downscale_input(image): def downscale_input(image):
samples = image.movedim(-1,1) samples = image.movedim(-1,1)
@ -38,13 +39,20 @@ def validate_and_cast_response (response):
# Get base64 image data # Get base64 image data
image_url = data[0].url image_url = data[0].url
if not image_url: b64_data = data[0].b64_json
raise Exception("No image URL was generated in the response") if not image_url and not b64_data:
raise Exception("No image was generated in the response")
if b64_data:
img_data = base64.b64decode(b64_data)
img = Image.open(io.BytesIO(img_data))
elif image_url:
img_response = requests.get(image_url) img_response = requests.get(image_url)
if img_response.status_code != 200: if img_response.status_code != 200:
raise Exception("Failed to download the image") raise Exception("Failed to download the image")
img = Image.open(io.BytesIO(img_response.content)) img = Image.open(io.BytesIO(img_response.content))
img = img.convert("RGB") # Ensure RGB format img = img.convert("RGB") # Ensure RGB format
# Convert to numpy array, normalize to float32 between 0 and 1 # Convert to numpy array, normalize to float32 between 0 and 1