feat: Support Imagen 4 Ingredients on Vertex

PiperOrigin-RevId: 808749595

Author: matthew29tang

google/genai/tests/models/test_edit_image.py

@@ -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',
@@ -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(

google/genai/types.py

@@ -6799,7 +6799,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."""
@@ -6816,7 +6816,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."""
@@ -13150,6 +13150,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."""