feat: support multi-image image generation

This commit is contained in:
QiuSW
2026-07-17 15:55:38 +08:00
parent bd274e26e3
commit 4a0becbb4d
23 changed files with 690 additions and 80 deletions
+2
View File
@@ -82,6 +82,7 @@ class ImageGenerationResult:
class MultimodalImage:
data: bytes
mime_type: str = "image/png"
filename: str = "image.png"
class Provider(Protocol):
@@ -108,6 +109,7 @@ class Provider(Protocol):
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
resolution: str = "1K",
aspect_ratio: str = "1:1",
parameters: Mapping[str, Any] | None = None,
+67 -9
View File
@@ -138,6 +138,7 @@ class ChatCompletionsProvider(BaseHttpProvider):
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
resolution: str = "1K",
aspect_ratio: str = "1:1",
parameters: Mapping[str, Any] | None = None,
@@ -149,6 +150,7 @@ class ChatCompletionsProvider(BaseHttpProvider):
prompt,
image=image,
image_mime_type=image_mime_type,
images=images,
parameters=parameters,
)
response = self.session.post(
@@ -243,6 +245,7 @@ class GeminiProvider(ChatCompletionsProvider):
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
resolution: str = "1K",
aspect_ratio: str = "1:1",
parameters: Mapping[str, Any] | None = None,
@@ -254,6 +257,7 @@ class GeminiProvider(ChatCompletionsProvider):
prompt,
image=image,
image_mime_type=image_mime_type,
images=images,
response_modalities=["TEXT", "IMAGE"],
parameters=parameters,
)
@@ -324,6 +328,7 @@ class ImagesGenerationProvider(BaseHttpProvider):
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
resolution: str = "1K",
aspect_ratio: str = "1:1",
parameters: Mapping[str, Any] | None = None,
@@ -337,8 +342,17 @@ class ImagesGenerationProvider(BaseHttpProvider):
"resolution": resolution,
"n": 1,
}
if image is not None:
payload["image_urls"] = [image_bytes_to_data_url(image, image_mime_type)]
image_inputs = generation_image_inputs(
image=image,
image_mime_type=image_mime_type,
image_filename=image_filename,
images=images,
)
if image_inputs:
payload["image_urls"] = [
image_bytes_to_data_url(item.data, item.mime_type)
for item in image_inputs
]
apply_extra_body(payload, model, parameters)
response = self.session.post(
url,
@@ -376,12 +390,19 @@ class ImagesEditsProvider(BaseHttpProvider):
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
resolution: str = "1K",
aspect_ratio: str = "1:1",
parameters: Mapping[str, Any] | None = None,
) -> ImageGenerationResult:
validate_model_config(model)
if image is None:
image_inputs = generation_image_inputs(
image=image,
image_mime_type=image_mime_type,
image_filename=image_filename,
images=images,
)
if not image_inputs:
raise AiCapabilityError("images edits provider requires an input image")
url = normalize_api_url(model.url, API_IMAGES_EDITS)
@@ -392,7 +413,15 @@ class ImagesEditsProvider(BaseHttpProvider):
"size": resolution_to_size(resolution),
}
apply_extra_body(data, model, parameters)
files = {"image": (image_filename, image, image_mime_type)}
files: dict[str, tuple[str, bytes, str]] | list[tuple[str, tuple[str, bytes, str]]]
if len(image_inputs) == 1:
item = image_inputs[0]
files = {"image": (item.filename, item.data, item.mime_type)}
else:
files = [
("image", (item.filename, item.data, item.mime_type))
for item in image_inputs
]
response = self.session.post(
url,
headers=self._headers(model),
@@ -443,13 +472,17 @@ def build_chat_image_payload(
*,
image: bytes | None = None,
image_mime_type: str = "image/png",
images: Sequence[MultimodalImage] | None = None,
parameters: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
return build_chat_text_payload(
return build_chat_vision_payload(
model,
prompt,
image=image,
image_mime_type=image_mime_type,
images=generation_image_inputs(
image=image,
image_mime_type=image_mime_type,
images=images,
),
parameters=parameters,
)
@@ -486,12 +519,17 @@ def build_gemini_payload(
*,
image: bytes | None = None,
image_mime_type: str = "image/png",
images: Sequence[MultimodalImage] | None = None,
response_modalities: list[str],
parameters: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
parts: list[dict[str, Any]] = [{"text": prompt}]
if image is not None:
data_url = image_bytes_to_data_url(image, image_mime_type)
for image_input in generation_image_inputs(
image=image,
image_mime_type=image_mime_type,
images=images,
):
data_url = image_bytes_to_data_url(image_input.data, image_input.mime_type)
mime_type, data = split_data_url(data_url)
parts.append({"inlineData": {"mimeType": mime_type, "data": data}})
payload: dict[str, Any] = {
@@ -522,6 +560,26 @@ def build_gemini_vision_payload(
return payload
def generation_image_inputs(
*,
image: bytes | None = None,
image_mime_type: str = "image/png",
image_filename: str = "image.png",
images: Sequence[MultimodalImage] | None = None,
) -> tuple[MultimodalImage, ...]:
if images is not None:
return tuple(images)
if image is None:
return ()
return (
MultimodalImage(
data=image,
mime_type=image_mime_type,
filename=image_filename,
),
)
def apply_extra_body(
payload: dict[str, Any],
model: ResolvedModel,