Skip to content

Commit e5e59fe

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
chore: Add support for background replacement to Imagen 3 Capabilities API
PiperOrigin-RevId: 708201423
1 parent d891b2c commit e5e59fe

File tree

1 file changed

+10
-2
lines changed

1 file changed

+10
-2
lines changed

vertexai/vision_models/_vision_models.py

+10-2
Original file line numberDiff line numberDiff line change
@@ -665,6 +665,7 @@ def _generate_images(
665665
"inpainting-remove",
666666
"outpainting",
667667
"product-image",
668+
"background-swap",
668669
"default",
669670
]
670671
] = None,
@@ -724,6 +725,8 @@ def _generate_images(
724725
mask)
725726
* product-image: Changes the background for the predominant
726727
product or subject in the image
728+
* background-swap: Changes the background for the predominant
729+
product or subject in the image
727730
* default: Default editing mode
728731
mask_mode: Solicits generation of the mask (v/s providing mask as an
729732
input). Supported values are:
@@ -884,7 +887,9 @@ class ID
884887
if isinstance(reference_image.config, ControlImageConfig):
885888
reference_image_instance["controlImageConfig"] = {
886889
"controlType": reference_image.config.control_type,
887-
"enableControlImageComputation": reference_image.config.enable_control_image_computation,
890+
"enableControlImageComputation": (
891+
reference_image.config.enable_control_image_computation
892+
),
888893
}
889894
shared_generation_parameters[
890895
f"reference_image_control_config_{reference_image.reference_id}"
@@ -952,6 +957,7 @@ class ID
952957
"inpainting-insert": "EDIT_MODE_INPAINT_INSERTION",
953958
"inpainting-remove": "EDIT_MODE_INPAINT_REMOVAL",
954959
"outpainting": "EDIT_MODE_OUTPAINT",
960+
"background-swap": "EDIT_MODE_BGSWAP",
955961
}
956962
capability_mode = (
957963
edit_mode_to_enum_map[edit_mode]
@@ -961,7 +967,9 @@ class ID
961967
parameters["editMode"] = capability_mode
962968
shared_generation_parameters["edit_mode"] = capability_mode
963969
else:
964-
edit_config["editMode"] = edit_mode
970+
edit_config["editMode"] = (
971+
edit_mode if edit_mode != "background-swap" else "inpainting-insert"
972+
)
965973
shared_generation_parameters["edit_mode"] = edit_mode
966974

967975
if mask is None and edit_mode is not None and edit_mode != "product-image":

0 commit comments

Comments
 (0)