Skip to content

Commit 0c498c5

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: Add safety filter levels, watermark support and person generation support for Imagen 2
Changelog: Changed the following parameters internally: - `editConfig` to `editConfigV6` - `segmentationClasses` to `classes` - `safetyFilterLevel` to `safetySetting` PiperOrigin-RevId: 619620158
1 parent 2e56acc commit 0c498c5

File tree

2 files changed

+5
-5
lines changed

2 files changed

+5
-5
lines changed

tests/unit/aiplatform/test_vision_models.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -499,7 +499,7 @@ def test_generate_images_gcs(self):
499499
actual_parameters["outputOptions"]["compressionQuality"]
500500
== compression_quality
501501
)
502-
assert actual_parameters["safetyFilterLevel"] == safety_filter_level
502+
assert actual_parameters["safetySetting"] == safety_filter_level
503503
assert actual_parameters["personGeneration"] == person_generation
504504

505505
assert len(image_response2.images) == number_of_images
@@ -639,7 +639,7 @@ def test_generate_images_requests_safety_filter_level(self):
639639
)
640640
predict_kwargs = mock_predict.call_args[1]
641641
actual_parameters = predict_kwargs["parameters"]
642-
assert actual_parameters["safetyFilterLevel"] == level
642+
assert actual_parameters["safetySetting"] == level
643643

644644
def test_generate_images_requests_person_generation(self):
645645
"""Tests that the model class generates person images."""

vertexai/vision_models/_vision_models.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -502,8 +502,8 @@ class ID
502502
shared_generation_parameters["mask_mode"] = mask_mode
503503

504504
if segmentation_classes is not None:
505-
parameters["editConfig"]["segmentationClasses"] = segmentation_classes
506-
shared_generation_parameters["segmentation_classes"] = segmentation_classes
505+
parameters["editConfig"]["classes"] = segmentation_classes
506+
shared_generation_parameters["classes"] = segmentation_classes
507507

508508
if mask_dilation is not None:
509509
parameters["editConfig"]["maskDilation"] = mask_dilation
@@ -527,7 +527,7 @@ class ID
527527
shared_generation_parameters["add_watermark"] = add_watermark
528528

529529
if safety_filter_level is not None:
530-
parameters["safetyFilterLevel"] = safety_filter_level
530+
parameters["safetySetting"] = safety_filter_level
531531
shared_generation_parameters["safety_filter_level"] = safety_filter_level
532532

533533
if person_generation is not None:

0 commit comments

Comments
 (0)