feat: LVM - Release ImageGenerationModel to GA

PiperOrigin-RevId: 658540981
googleapis · Aug 1, 2024 · 718c199 · 718c199
1 parent 71c0fd3
commit 718c199
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/tests/unit/aiplatform/test_vision_models.py b/tests/unit/aiplatform/test_vision_models.py
@@ -203,7 +203,7 @@ def teardown_method(self):
 
     def _get_image_generation_model(
         self,
-    ) -> preview_vision_models.ImageGenerationModel:
+    ) -> ga_vision_models.ImageGenerationModel:
         """Gets the image generation model."""
         aiplatform.init(
             project=_TEST_PROJECT,
@@ -216,7 +216,7 @@ def _get_image_generation_model(
                 _IMAGE_GENERATION_PUBLISHER_MODEL_DICT
             ),
         ) as mock_get_publisher_model:
-            model = preview_vision_models.ImageGenerationModel.from_pretrained(
+            model = ga_vision_models.ImageGenerationModel.from_pretrained(
                 "imagegeneration@002"
             )
 

diff --git a/vertexai/vision_models/__init__.py b/vertexai/vision_models/__init__.py
@@ -15,8 +15,11 @@
 """Classes for working with vision models."""
 
 from vertexai.vision_models._vision_models import (
+    GeneratedImage,
     Image,
     ImageCaptioningModel,
+    ImageGenerationModel,
+    ImageGenerationResponse,
     ImageQnAModel,
     ImageTextModel,
     MultiModalEmbeddingModel,
@@ -27,8 +30,11 @@
 )
 
 __all__ = [
+    "GeneratedImage",
     "Image",
     "ImageCaptioningModel",
+    "ImageGenerationModel",
+    "ImageGenerationResponse",
     "ImageQnAModel",
     "ImageTextModel",
     "MultiModalEmbeddingModel",