123 lines
4.1 KiB
Python
123 lines
4.1 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
from typing import List, Union
|
|
|
|
|
|
MIN_IMAGE_PIXELS = 655_360
|
|
MAX_IMAGE_PIXELS = 8_294_400
|
|
MAX_IMAGE_EDGE = 3840
|
|
IMAGE_DIMENSION_STEP = 16
|
|
MAX_IMAGE_ASPECT_RATIO = 3.0
|
|
|
|
|
|
def validate_image_size(width: int, height: int):
|
|
"""Validate GPT-Image-2 custom output size constraints."""
|
|
width = int(width)
|
|
height = int(height)
|
|
pixels = width * height
|
|
long_edge = max(width, height)
|
|
short_edge = min(width, height)
|
|
|
|
if width <= 0 or height <= 0:
|
|
return None, "Width and height must be positive."
|
|
if width % IMAGE_DIMENSION_STEP != 0 or height % IMAGE_DIMENSION_STEP != 0:
|
|
return None, f"Width and height must be multiples of {IMAGE_DIMENSION_STEP}."
|
|
if long_edge > MAX_IMAGE_EDGE:
|
|
return None, f"Maximum edge length is {MAX_IMAGE_EDGE}px."
|
|
if pixels < MIN_IMAGE_PIXELS:
|
|
return None, f"Total pixels must be at least {MIN_IMAGE_PIXELS}."
|
|
if pixels > MAX_IMAGE_PIXELS:
|
|
return None, f"Total pixels must be at most {MAX_IMAGE_PIXELS}."
|
|
if long_edge / short_edge > MAX_IMAGE_ASPECT_RATIO:
|
|
return None, f"Aspect ratio must be no more than {MAX_IMAGE_ASPECT_RATIO:g}:1."
|
|
return f"{width}x{height}", None
|
|
|
|
|
|
def resolve_model_choice(choice: str, custom_value: str = "", omit_values=None) -> str:
|
|
"""Resolve dropdown/custom model fields; return empty string to omit a model."""
|
|
omit_values = set(omit_values or [])
|
|
choice = (choice or "").strip()
|
|
custom_value = (custom_value or "").strip()
|
|
if choice == "custom":
|
|
return custom_value
|
|
if choice in omit_values:
|
|
return ""
|
|
return choice
|
|
|
|
|
|
def tensor_to_pil(image: torch.Tensor) -> List[Image.Image]:
|
|
"""ComfyUI tensor [B,H,W,C] or [H,W,C] -> list of PIL Images."""
|
|
if image is None:
|
|
return []
|
|
arr = image.cpu().numpy() if hasattr(image, "cpu") else np.array(image)
|
|
if len(arr.shape) == 3:
|
|
arr = arr[np.newaxis, ...]
|
|
images = []
|
|
for i in range(arr.shape[0]):
|
|
img = arr[i]
|
|
# Normalize: float [0,1] -> uint8 [0,255]; int values pass through clipped
|
|
if img.dtype == np.float32 or img.dtype == np.float64:
|
|
img = np.clip(img * 255.0, 0, 255).astype(np.uint8)
|
|
else:
|
|
img = np.clip(img, 0, 255).astype(np.uint8)
|
|
if len(img.shape) == 2:
|
|
pil_img = Image.fromarray(img, mode="L")
|
|
elif img.shape[2] == 4:
|
|
pil_img = Image.fromarray(img, mode="RGBA")
|
|
else:
|
|
pil_img = Image.fromarray(img, mode="RGB")
|
|
images.append(pil_img)
|
|
return images
|
|
|
|
|
|
def pil_to_tensor(images: Union[Image.Image, List[Image.Image]]) -> torch.Tensor:
|
|
"""PIL or list of PIL -> ComfyUI tensor [B,H,W,C], value range [0,1]."""
|
|
if isinstance(images, Image.Image):
|
|
images = [images]
|
|
tensors = []
|
|
for img in images:
|
|
if img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
arr = np.array(img).astype(np.float32) / 255.0
|
|
tensors.append(torch.from_numpy(arr))
|
|
return torch.stack(tensors)
|
|
|
|
|
|
def pil_to_bytes(img: Image.Image, format: str = "PNG") -> BytesIO:
|
|
"""Convert a single PIL image to BytesIO for API upload."""
|
|
buf = BytesIO()
|
|
img.save(buf, format=format)
|
|
buf.seek(0)
|
|
return buf
|
|
|
|
|
|
def tensor_to_bytes_list(image: torch.Tensor) -> List[BytesIO]:
|
|
"""Convert tensor batch to list of BytesIO objects."""
|
|
pils = tensor_to_pil(image)
|
|
return [pil_to_bytes(p) for p in pils]
|
|
|
|
|
|
def format_openai_error(e: Exception) -> str:
|
|
"""Extract a user-friendly message from an OpenAI SDK exception."""
|
|
body = getattr(e, "body", None)
|
|
if isinstance(body, dict):
|
|
err = body.get("error", {})
|
|
msg = err.get("message", "")
|
|
code = err.get("code", "")
|
|
etype = err.get("type", "")
|
|
parts = []
|
|
if code:
|
|
parts.append(f"[{code}]")
|
|
if msg:
|
|
parts.append(msg)
|
|
if etype and etype != code:
|
|
parts.append(f"({etype})")
|
|
if parts:
|
|
return " ".join(parts)
|
|
msg = getattr(e, "message", None)
|
|
if msg:
|
|
return str(msg)
|
|
return str(e)
|