feat: support multi-image image generation
This commit is contained in:
QiuSW
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user