Basically works.

Change-Id: I28364ab70b2a263b29026f2cf2d1d4f807d88f53

Author: MarkDaoust

google/generativeai/__init__.py

@@ -59,6 +59,8 @@
 from google.generativeai.generative_models import GenerativeModel
 from google.generativeai.generative_models import ChatSession
 
+from google.generativeai.vision_models import *
+
 from google.generativeai.models import list_models
 from google.generativeai.models import list_tuned_models
 

google/generativeai/client.py

@@ -377,3 +377,11 @@ def get_default_permission_client() -> glm.PermissionServiceClient:
 
 def get_default_permission_async_client() -> glm.PermissionServiceAsyncClient:
     return _client_manager.get_default_client("permission_async")
+
+
+def get_default_prediction_client() -> glm.PermissionServiceClient:
+    return _client_manager.get_default_client("prediction")
+
+
+def get_default_prediction_async_client() -> glm.PermissionServiceAsyncClient:
+    return _client_manager.get_default_client("prediction_async")

google/generativeai/vision_models/_vision_models.py

@@ -25,6 +25,9 @@
 import typing
 from typing import Any, Dict, List, Literal, Optional, Union
 
+from google.generativeai import client
+from google.generativeai import protos
+
 from google.protobuf import struct_pb2
 
 from proto.marshal.collections import maps
@@ -43,6 +46,7 @@
     PIL_Image = None
 
 
+# This is to get around https://github.com/googleapis/proto-plus-python/issues/488
 def to_value(value) -> struct_pb2.Value:
     """Return a protobuf Value object representing this value."""
     if isinstance(value, struct_pb2.Value):
@@ -61,6 +65,7 @@ def to_value(value) -> struct_pb2.Value:
         return struct_pb2.Value(struct_value=to_mapping_value(value))
     raise ValueError("Unable to coerce value: %r" % value)
 
+
 def to_list_value(value) -> struct_pb2.ListValue:
     # We got a proto, or else something we sent originally.
     # Preserve the instance we have.
@@ -70,9 +75,8 @@ def to_list_value(value) -> struct_pb2.ListValue:
         return struct_pb2.ListValue(values=[v for v in value.pb])
 
     # We got a list (or something list-like); convert it.
-    return struct_pb2.ListValue(
-        values=[to_value(v) for v in value]
-    )
+    return struct_pb2.ListValue(values=[to_value(v) for v in value])
+
 
 def to_mapping_value(value) -> struct_pb2.Struct:
     # We got a proto, or else something we sent originally.
@@ -85,12 +89,7 @@ def to_mapping_value(value) -> struct_pb2.Struct:
         )
 
     # We got a dict (or something dict-like); convert it.
-    return struct_pb2.Struct(
-        fields={
-            k: to_value(v) for k, v in value.items()
-        }
-    )
-
+    return struct_pb2.Struct(fields={k: to_value(v) for k, v in value.items()})
 
 
 _SUPPORTED_UPSCALING_SIZES = [2048, 4096]
@@ -131,7 +130,6 @@ def load_from_file(location: str) -> "Image":
         image = Image(image_bytes=image_bytes)
         return image
 
-
     @property
     def _image_bytes(self) -> bytes:
         return self._loaded_bytes
@@ -206,9 +204,16 @@ class ImageGenerationModel:
         response[0].save("image1.png")
     """
 
-    __module__ = "vertexai.preview.vision_models"
+    def __init__(self, model_id: str):
+        if not model_id.startswith("models"):
+            model_id = f"models/{model_id}"
+        self.model_name = model_id
+        self._client = None
 
-    _INSTANCE_SCHEMA_URI = "gs://google-cloud-aiplatform/schema/predict/instance/vision_generative_model_1.0.0.yaml"
+    @classmethod
+    def from_pretrained(cls, model_name: str):
+        """For vertex compatibility"""
+        return cls(model_name)
 
     def _generate_images(
         self,
@@ -242,9 +247,7 @@ def _generate_images(
         safety_filter_level: Optional[
             Literal["block_most", "block_some", "block_few", "block_fewest"]
         ] = None,
-        person_generation: Optional[
-            Literal["dont_allow", "allow_adult", "allow_all"]
-        ] = None,
+        person_generation: Optional[Literal["dont_allow", "allow_adult", "allow_all"]] = None,
     ) -> "ImageGenerationResponse":
         """Generates images from text prompt.
 
@@ -312,6 +315,8 @@ class ID
         Returns:
             An `ImageGenerationResponse` object.
         """
+        if self._client is None:
+            self._client = client.get_default_prediction_client()
         # Note: Only a single prompt is supported by the service.
         instance = {"prompt": prompt}
         shared_generation_parameters = {
@@ -412,11 +417,14 @@ class ID
             parameters["personGeneration"] = person_generation
             shared_generation_parameters["person_generation"] = person_generation
 
-        response = self._endpoint.predict(
-            instances=[to_value(instance)],
-            parameters=parameters,
+        # This is to get around https://github.com/googleapis/proto-plus-python/issues/488
+        pr = protos.PredictRequest.pb()
+        request = pr(
+            model=self.model_name, instances=[to_value(instance)], parameters=to_value(parameters)
         )
 
+        response = self._client.predict(request)
+
         generated_images: List["GeneratedImage"] = []
         for idx, prediction in enumerate(response.predictions):
             generation_parameters = dict(shared_generation_parameters)
@@ -444,9 +452,7 @@ def generate_images(
         safety_filter_level: Optional[
             Literal["block_most", "block_some", "block_few", "block_fewest"]
         ] = None,
-        person_generation: Optional[
-            Literal["dont_allow", "allow_adult", "allow_all"]
-        ] = None,
+        person_generation: Optional[Literal["dont_allow", "allow_adult", "allow_all"]] = None,
     ) -> "ImageGenerationResponse":
         """Generates images from text prompt.
 
@@ -510,9 +516,7 @@ def edit_image(
         number_of_images: int = 1,
         guidance_scale: Optional[float] = None,
         edit_mode: Optional[
-            Literal[
-                "inpainting-insert", "inpainting-remove", "outpainting", "product-image"
-            ]
+            Literal["inpainting-insert", "inpainting-remove", "outpainting", "product-image"]
         ] = None,
         mask_mode: Optional[Literal["background", "foreground", "semantic"]] = None,
         segmentation_classes: Optional[List[str]] = None,
@@ -525,9 +529,7 @@ def edit_image(
         safety_filter_level: Optional[
             Literal["block_most", "block_some", "block_few", "block_fewest"]
         ] = None,
-        person_generation: Optional[
-            Literal["dont_allow", "allow_adult", "allow_all"]
-        ] = None,
+        person_generation: Optional[Literal["dont_allow", "allow_adult", "allow_all"]] = None,
     ) -> "ImageGenerationResponse":
         """Edits an existing image based on text prompt.
 
@@ -717,9 +719,7 @@ def upscale_image(
 
         parameters["outputOptions"] = {"mimeType": output_mime_type}
         if output_mime_type == "image/jpeg" and output_compression_quality is not None:
-            parameters["outputOptions"][
-                "compressionQuality"
-            ] = output_compression_quality
+            parameters["outputOptions"]["compressionQuality"] = output_compression_quality
 
         response = self._endpoint.predict(
             instances=[to_value(instance)],
@@ -825,9 +825,7 @@ def save(self, location: str, include_generation_parameters: bool = True):
             if not self._generation_parameters:
                 raise ValueError("Image does not have generation parameters.")
             if not PIL_Image:
-                raise ValueError(
-                    "The PIL module is required for saving generation parameters."
-                )
+                raise ValueError("The PIL module is required for saving generation parameters.")
 
             exif = self._pil_image.getexif()
             exif[_EXIF_USER_COMMENT_TAG_IDX] = json.dumps(
@@ -836,4 +834,3 @@ def save(self, location: str, include_generation_parameters: bool = True):
             self._pil_image.save(location, exif=exif)
         else:
             super().save(location=location)
-