fix: LVM - Fixed negative prompt in ImageGenerationModel

PiperOrigin-RevId: 574652096
googleapis · Oct 19, 2023 · cbe3a0d · cbe3a0d
1 parent d9ced10
commit cbe3a0d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/tests/unit/aiplatform/test_vision_models.py b/tests/unit/aiplatform/test_vision_models.py
@@ -251,7 +251,7 @@ def test_generate_images(self):
             actual_parameters = predict_kwargs["parameters"]
             actual_instance = predict_kwargs["instances"][0]
             assert actual_instance["prompt"] == prompt1
-            assert actual_instance["negativePrompt"] == negative_prompt1
+            assert actual_parameters["negativePrompt"] == negative_prompt1
             # TODO(b/295946075) The service stopped supporting image sizes.
             # assert actual_parameters["sampleImageSize"] == str(max(width, height))
             # assert actual_parameters["aspectRatio"] == f"{width}:{height}"

diff --git a/vertexai/vision_models/_vision_models.py b/vertexai/vision_models/_vision_models.py
@@ -174,10 +174,6 @@ def _generate_images(
             "number_of_images_in_batch": number_of_images,
         }
 
-        if negative_prompt:
-            instance["negativePrompt"] = negative_prompt
-            shared_generation_parameters["negative_prompt"] = negative_prompt
-
         if base_image:
             base_image_base64 = (
                 base_image._as_base64_string()
@@ -207,6 +203,9 @@ def _generate_images(
                 parameters["aspectRatio"] = f"{width}:{height}"
 
         parameters["sampleCount"] = number_of_images
+        if negative_prompt:
+            parameters["negativePrompt"] = negative_prompt
+            shared_generation_parameters["negative_prompt"] = negative_prompt
 
         if seed is not None:
             # Note: String seed and numerical seed give different results