feat: Add safety_filter_level and person_generation for Imagen upscaling

PiperOrigin-RevId: 819915543

Author: matthew29tang

google/genai/models.py

@@ -3516,6 +3516,20 @@ def _UpscaleImageAPIConfig_to_vertex(
         getv(from_object, ['output_gcs_uri']),
     )
 
+  if getv(from_object, ['safety_filter_level']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'safetySetting'],
+        getv(from_object, ['safety_filter_level']),
+    )
+
+  if getv(from_object, ['person_generation']) is not None:
+    setv(
+        parent_object,
+        ['parameters', 'personGeneration'],
+        getv(from_object, ['person_generation']),
+    )
+
   if getv(from_object, ['include_rai_reason']) is not None:
     setv(
         parent_object,
@@ -5376,6 +5390,8 @@ def upscale_image(
     api_config = types._UpscaleImageAPIConfigDict(
         http_options=config_dct.get('http_options', None),
         output_gcs_uri=config_dct.get('output_gcs_uri', None),
+        safety_filter_level=config_dct.get('safety_filter_level', None),
+        person_generation=config_dct.get('person_generation', None),
         include_rai_reason=config_dct.get('include_rai_reason', None),
         output_mime_type=config_dct.get('output_mime_type', None),
         output_compression_quality=config_dct.get(
@@ -7200,6 +7216,8 @@ async def upscale_image(
     api_config = types._UpscaleImageAPIConfigDict(
         http_options=config_dct.get('http_options', None),
         output_gcs_uri=config_dct.get('output_gcs_uri', None),
+        safety_filter_level=config_dct.get('safety_filter_level', None),
+        person_generation=config_dct.get('person_generation', None),
         include_rai_reason=config_dct.get('include_rai_reason', None),
         output_mime_type=config_dct.get('output_mime_type', None),
         output_compression_quality=config_dct.get(

google/genai/tests/models/test_upscale_image.py

@@ -49,6 +49,8 @@
             upscale_factor='x2',
             config=types.UpscaleImageConfig(
                 include_rai_reason=True,
+                person_generation=types.PersonGeneration.ALLOW_ADULT,
+                safety_filter_level=types.SafetyFilterLevel.BLOCK_LOW_AND_ABOVE,
                 output_mime_type='image/jpeg',
                 output_compression_quality=80,
                 enhance_input_image=True,
@@ -107,6 +109,8 @@ async def test_upscale_async(client):
         image=types.Image.from_file(location=IMAGE_FILE_PATH),
         upscale_factor='x2',
         config=types.UpscaleImageConfig(
+            person_generation=types.PersonGeneration.ALLOW_ADULT,
+            safety_filter_level=types.SafetyFilterLevel.BLOCK_LOW_AND_ABOVE,
             include_rai_reason=True,
             output_mime_type='image/jpeg',
             output_compression_quality=80,

google/genai/types.py

@@ -7384,6 +7384,12 @@ class _UpscaleImageAPIConfig(_common.BaseModel):
       default=None,
       description="""Cloud Storage URI used to store the generated images.""",
   )
+  safety_filter_level: Optional[SafetyFilterLevel] = Field(
+      default=None, description="""Filter level for safety filtering."""
+  )
+  person_generation: Optional[PersonGeneration] = Field(
+      default=None, description="""Allows generation of people by the model."""
+  )
   include_rai_reason: Optional[bool] = Field(
       default=None,
       description="""Whether to include a reason for filtered-out images in the
@@ -7432,6 +7438,12 @@ class _UpscaleImageAPIConfigDict(TypedDict, total=False):
   output_gcs_uri: Optional[str]
   """Cloud Storage URI used to store the generated images."""
 
+  safety_filter_level: Optional[SafetyFilterLevel]
+  """Filter level for safety filtering."""
+
+  person_generation: Optional[PersonGeneration]
+  """Allows generation of people by the model."""
+
   include_rai_reason: Optional[bool]
   """Whether to include a reason for filtered-out images in the
       response."""
@@ -13546,6 +13558,12 @@ class UpscaleImageConfig(_common.BaseModel):
       default=None,
       description="""Cloud Storage URI used to store the generated images.""",
   )
+  safety_filter_level: Optional[SafetyFilterLevel] = Field(
+      default=None, description="""Filter level for safety filtering."""
+  )
+  person_generation: Optional[PersonGeneration] = Field(
+      default=None, description="""Allows generation of people by the model."""
+  )
   include_rai_reason: Optional[bool] = Field(
       default=None,
       description="""Whether to include a reason for filtered-out images in the
@@ -13593,6 +13611,12 @@ class UpscaleImageConfigDict(TypedDict, total=False):
   output_gcs_uri: Optional[str]
   """Cloud Storage URI used to store the generated images."""
 
+  safety_filter_level: Optional[SafetyFilterLevel]
+  """Filter level for safety filtering."""
+
+  person_generation: Optional[PersonGeneration]
+  """Allows generation of people by the model."""
+
   include_rai_reason: Optional[bool]
   """Whether to include a reason for filtered-out images in the
       response."""