Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions google/genai/tests/models/test_edit_image.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,23 @@
),
)

dog_content_ref_image = types.ContentReferenceImage(
reference_id=1,
reference_image=types.Image(
gcs_uri='gs://genai-sdk-tests/inputs/images/dog.jpg'
),
)

cyberpunk_style_ref_image = types.StyleReferenceImage(
reference_id=2,
reference_image=types.Image(
gcs_uri='gs://genai-sdk-tests/inputs/images/cyberpunk.jpg'
),
config=types.StyleReferenceConfig(
style_description='cyberpunk style',
),
)

test_table: list[pytest_helper.TestTableItem] = [
pytest_helper.TestTableItem(
name='test_edit_mask_inpaint_insert',
Expand Down Expand Up @@ -176,6 +193,25 @@
),
),
),
pytest_helper.TestTableItem(
name='test_edit_content_image_ingredients',
exception_if_mldev='only supported in the Vertex AI client',
parameters=types._EditImageParameters(
model='imagen-4.0-ingredients-preview',
prompt=(
'Dog in [1] sleeping on the ground at the bottom of the image'
' with the cyberpunk city landscape in [2] in the background'
' visible on the side of the mug.'
),
reference_images=[dog_content_ref_image, cyberpunk_style_ref_image],
config=types.EditImageConfig(
number_of_images=1,
aspect_ratio='9:16',
include_rai_reason=True,
output_mime_type='image/jpeg',
),
),
),
]

pytestmark = pytest_helper.setup(
Expand Down
56 changes: 54 additions & 2 deletions google/genai/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -6786,7 +6786,7 @@ class _EditImageParameters(_common.BaseModel):
description="""A text description of the edit to apply to the image.""",
)
reference_images: Optional[list[_ReferenceImageAPI]] = Field(
default=None, description="""The reference images for Imagen 3 editing."""
default=None, description="""The reference images for editing."""
)
config: Optional[EditImageConfig] = Field(
default=None, description="""Configuration for editing."""
Expand All @@ -6803,7 +6803,7 @@ class _EditImageParametersDict(TypedDict, total=False):
"""A text description of the edit to apply to the image."""

reference_images: Optional[list[_ReferenceImageAPIDict]]
"""The reference images for Imagen 3 editing."""
"""The reference images for editing."""

config: Optional[EditImageConfigDict]
"""Configuration for editing."""
Expand Down Expand Up @@ -13137,6 +13137,58 @@ class SubjectReferenceImageDict(TypedDict, total=False):
]


class ContentReferenceImage(_common.BaseModel):
"""A content reference image.

A content reference image represents a subject to reference (ex. person,
product, animal) provided by the user. It can optionally be provided in
addition to a style reference image (ex. background, style reference).
"""

reference_image: Optional[Image] = Field(
default=None,
description="""The reference image for the editing operation.""",
)
reference_id: Optional[int] = Field(
default=None, description="""The id of the reference image."""
)
reference_type: Optional[str] = Field(
default=None,
description="""The type of the reference image. Only set by the SDK.""",
)

@pydantic.model_validator(mode='before')
@classmethod
def _validate_mask_image_config(self, values: Any) -> Any:
if 'reference_type' in values:
raise ValueError('Cannot set internal reference_type field directly.')
values['reference_type'] = 'REFERENCE_TYPE_CONTENT'
return values


class ContentReferenceImageDict(TypedDict, total=False):
"""A content reference image.

A content reference image represents a subject to reference (ex. person,
product, animal) provided by the user. It can optionally be provided in
addition to a style reference image (ex. background, style reference).
"""

reference_image: Optional[ImageDict]
"""The reference image for the editing operation."""

reference_id: Optional[int]
"""The id of the reference image."""

reference_type: Optional[str]
"""The type of the reference image. Only set by the SDK."""


ContentReferenceImageOrDict = Union[
ContentReferenceImage, ContentReferenceImageDict
]


class LiveServerSetupComplete(_common.BaseModel):
"""Sent in response to a `LiveGenerateContentSetup` message from the client."""

Expand Down