Image gen

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

@@ -384,3 +384,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/generative_models.py

@@ -4,7 +4,7 @@
 
 from collections.abc import Iterable
 import textwrap
-from typing import Any, Union, overload
+from typing import Any, Literal, Union, overload
 import reprlib
 
 # pylint: disable=bad-continuation, line-too-long

google/generativeai/vision_models/__init__.py

@@ -0,0 +1,29 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Classes for working with vision models."""
+
+from google.generativeai.vision_models._vision_models import (
+    Image,
+    GeneratedImage,
+    ImageGenerationModel,
+    ImageGenerationResponse,
+)
+
+__all__ = [
+    "Image",
+    "GeneratedImage",
+    "ImageGenerationModel",
+    "ImageGenerationResponse",
+]