phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor vertex to use spearate image gen folder

Browse source
This commit is contained in:
Ishaan Jaff 2024-09-02 17:34:43 -07:00
parent a9c9967b6d
commit e779739f72
1 changed files with 5 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/main.py
View file

@ -126,6 +126,9 @@ from .llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gem
from .llms.vertex_ai_and_google_ai_studio.gemini_embeddings.batch_embed_content_handler import ( from .llms.vertex_ai_and_google_ai_studio.gemini_embeddings.batch_embed_content_handler import (
GoogleBatchEmbeddings, GoogleBatchEmbeddings,
) )
from .llms.vertex_ai_and_google_ai_studio.image_generation.image_generation_handler import (
VertexImageGeneration,
)
from .llms.vertex_ai_and_google_ai_studio.multimodal_embeddings.embedding_handler import ( from .llms.vertex_ai_and_google_ai_studio.multimodal_embeddings.embedding_handler import (
VertexMultimodalEmbedding, VertexMultimodalEmbedding,
) )
@ -180,6 +183,7 @@ bedrock_converse_chat_completion = BedrockConverseLLM()
bedrock_embedding = BedrockEmbedding() bedrock_embedding = BedrockEmbedding()
vertex_chat_completion = VertexLLM() vertex_chat_completion = VertexLLM()
vertex_multimodal_embedding = VertexMultimodalEmbedding() vertex_multimodal_embedding = VertexMultimodalEmbedding()
vertex_image_generation = VertexImageGeneration()
google_batch_embeddings = GoogleBatchEmbeddings() google_batch_embeddings = GoogleBatchEmbeddings()
vertex_partner_models_chat_completion = VertexAIPartnerModels() vertex_partner_models_chat_completion = VertexAIPartnerModels()
vertex_text_to_speech = VertexTextToSpeechAPI() vertex_text_to_speech = VertexTextToSpeechAPI()
@ -4534,7 +4538,7 @@ def image_generation(
or optional_params.pop("vertex_ai_credentials", None) or optional_params.pop("vertex_ai_credentials", None)
or get_secret("VERTEXAI_CREDENTIALS") or get_secret("VERTEXAI_CREDENTIALS")
) )
model_response = vertex_chat_completion.image_generation( model_response = vertex_image_generation.image_generation(
model=model, model=model,
prompt=prompt, prompt=prompt,
timeout=timeout, timeout=timeout,