Spaces:
Running
Running
File size: 3,040 Bytes
c47ec10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | from __future__ import annotations
from io import BytesIO
from typing import Any, Iterator
from PIL import Image
from services.protocol.conversation import (
ConversationRequest,
ImageGenerationError,
collect_image_outputs,
count_text_tokens,
encode_images,
stream_image_chunks,
stream_image_outputs_with_pool,
)
from utils.image_tokens import count_image_inputs_tokens, count_image_output_items_tokens, image_usage
def _composite_mask(
images: list[tuple[bytes, str, str]],
masks: list[tuple[bytes, str, str]],
) -> list[tuple[bytes, str, str]]:
"""ๅฐ mask ็ alpha ้้ๅๆๅฐๅพ็ไธญ๏ผๆ ่ฏ้่ฆ็ผ่พ็ๅบๅใ
mask ็้ๆๅบๅ๏ผไฝ alpha๏ผ= ้่ฆ็ผ่พ็ๅบๅ๏ผ
mask ็ไธ้ๆๅบๅ๏ผ้ซ alpha๏ผ= ไฟ็็ๅบๅใ
ๅฆๆๆ mask ๅ่ฟๅๅๅพใ
"""
if not masks:
return images
result: list[tuple[bytes, str, str]] = []
for i, (data, filename, mime_type) in enumerate(images):
mask_data = masks[i][0] if i < len(masks) else masks[-1][0]
img = Image.open(BytesIO(data)).convert("RGBA")
mask_img = Image.open(BytesIO(mask_data))
if mask_img.mode == "RGBA":
alpha = mask_img.split()[3]
elif mask_img.mode == "L":
alpha = mask_img
else:
alpha = mask_img.convert("L")
alpha = alpha.resize(img.size, Image.LANCZOS)
img.putalpha(alpha)
buf = BytesIO()
img.save(buf, format="PNG")
result.append((buf.getvalue(), filename, "image/png"))
return result
def handle(body: dict[str, Any]) -> dict[str, Any] | Iterator[dict[str, Any]]:
prompt = str(body.get("prompt") or "")
images = body.get("images") or []
masks = body.get("mask") or []
images = _composite_mask(images, masks)
model = str(body.get("model") or "gpt-image-2")
n = int(body.get("n") or 1)
size = body.get("size")
quality = str(body.get("quality") or "auto")
response_format = str(body.get("response_format") or "b64_json")
base_url = str(body.get("base_url") or "") or None
progress_callback = body.get("progress_callback")
encoded_images = encode_images(images)
if not encoded_images:
raise ImageGenerationError("image is required")
outputs = stream_image_outputs_with_pool(ConversationRequest(
prompt=prompt,
model=model,
n=n,
size=size,
quality=quality,
response_format=response_format,
base_url=base_url,
images=encoded_images,
message_as_error=True,
progress_callback=progress_callback,
))
if body.get("stream"):
return stream_image_chunks(outputs)
result = collect_image_outputs(outputs)
result["usage"] = image_usage(
input_text_tokens=count_text_tokens(prompt, model),
input_image_tokens=count_image_inputs_tokens(images, model),
output_tokens=count_image_output_items_tokens(result.get("data"), size, quality),
)
return result
|