feat: Add enhance_input_image and image_preservation_factor fields for upscale_image

PiperOrigin-RevId: 775438709

Author: matthew29tang

google/genai/models.py

@@ -2862,6 +2862,20 @@ def _UpscaleImageAPIConfig_to_vertex(
         getv(from_object, ['output_compression_quality']),
     )
 
+  if getv(from_object, ['enhance_input_image']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'upscaleConfig', 'enhanceInputImage'],
+        getv(from_object, ['enhance_input_image']),
+    )
+
+  if getv(from_object, ['image_preservation_factor']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'upscaleConfig', 'imagePreservationFactor'],
+        getv(from_object, ['image_preservation_factor']),
+    )
+
   if getv(from_object, ['number_of_images']) is not None:
     setv(
         parent_object,
@@ -6240,6 +6254,10 @@ def upscale_image(
         output_compression_quality=config_dct.get(
             'output_compression_quality', None
         ),
+        enhance_input_image=config_dct.get('enhance_input_image', None),
+        image_preservation_factor=config_dct.get(
+            'image_preservation_factor', None
+        ),
     )  # pylint: disable=protected-access
 
     # Provide default values through API config.
@@ -7832,6 +7850,10 @@ async def upscale_image(
         output_compression_quality=config_dct.get(
             'output_compression_quality', None
         ),
+        enhance_input_image=config_dct.get('enhance_input_image', None),
+        image_preservation_factor=config_dct.get(
+            'image_preservation_factor', None
+        ),
     )  # pylint: disable=protected-access
 
     # Provide default values through API config.

google/genai/tests/models/test_upscale_image.py

@@ -50,6 +50,8 @@
                 'include_rai_reason': True,
                 'output_mime_type': 'image/jpeg',
                 'output_compression_quality': 80,
+                'enhance_input_image': True,
+                'image_preservation_factor': 0.6,
             },
         ),
     ),
@@ -94,6 +96,8 @@ async def test_upscale_async(client):
             'include_rai_reason': True,
             'output_mime_type': 'image/jpeg',
             'output_compression_quality': 80,
+            'enhance_input_image': True,
+            'image_preservation_factor': 0.6,
         },
     )
     assert response.generated_images[0].image.image_bytes

google/genai/types.py

@@ -6282,6 +6282,19 @@ class _UpscaleImageAPIConfig(_common.BaseModel):
       description="""The level of compression if the ``output_mime_type`` is
       ``image/jpeg``.""",
   )
+  enhance_input_image: Optional[bool] = Field(
+      default=None,
+      description="""Whether to add an image enhancing step before upscaling.
+      It is expected to suppress the noise and JPEG compression artifacts
+      from the input image.""",
+  )
+  image_preservation_factor: Optional[float] = Field(
+      default=None,
+      description="""With a higher image preservation factor, the original image
+      pixels are more respected. With a lower image preservation factor, the
+      output image will have be more different from the input image, but
+      with finer details and less noise.""",
+  )
   number_of_images: Optional[int] = Field(default=None, description="""""")
   mode: Optional[str] = Field(default=None, description="""""")
 
@@ -6307,6 +6320,17 @@ class _UpscaleImageAPIConfigDict(TypedDict, total=False):
   """The level of compression if the ``output_mime_type`` is
       ``image/jpeg``."""
 
+  enhance_input_image: Optional[bool]
+  """Whether to add an image enhancing step before upscaling.
+      It is expected to suppress the noise and JPEG compression artifacts
+      from the input image."""
+
+  image_preservation_factor: Optional[float]
+  """With a higher image preservation factor, the original image
+      pixels are more respected. With a lower image preservation factor, the
+      output image will have be more different from the input image, but
+      with finer details and less noise."""
+
   number_of_images: Optional[int]
   """"""
 
@@ -10518,6 +10542,19 @@ class UpscaleImageConfig(_common.BaseModel):
       description="""The level of compression if the ``output_mime_type`` is
       ``image/jpeg``.""",
   )
+  enhance_input_image: Optional[bool] = Field(
+      default=None,
+      description="""Whether to add an image enhancing step before upscaling.
+      It is expected to suppress the noise and JPEG compression artifacts
+      from the input image.""",
+  )
+  image_preservation_factor: Optional[float] = Field(
+      default=None,
+      description="""With a higher image preservation factor, the original image
+      pixels are more respected. With a lower image preservation factor, the
+      output image will have be more different from the input image, but
+      with finer details and less noise.""",
+  )
 
 
 class UpscaleImageConfigDict(TypedDict, total=False):
@@ -10542,6 +10579,17 @@ class UpscaleImageConfigDict(TypedDict, total=False):
   """The level of compression if the ``output_mime_type`` is
       ``image/jpeg``."""
 
+  enhance_input_image: Optional[bool]
+  """Whether to add an image enhancing step before upscaling.
+      It is expected to suppress the noise and JPEG compression artifacts
+      from the input image."""
+
+  image_preservation_factor: Optional[float]
+  """With a higher image preservation factor, the original image
+      pixels are more respected. With a lower image preservation factor, the
+      output image will have be more different from the input image, but
+      with finer details and less noise."""
+
 
 UpscaleImageConfigOrDict = Union[UpscaleImageConfig, UpscaleImageConfigDict]