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

feat(main.py): add support for image generation endpoint

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-16 21:07:29 -08:00
parent e62327dd92
commit 51cb16a015
7 changed files with 366 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

37
litellm/tests/test_image_generation.py Normal file
View file

@ -0,0 +1,37 @@
# What this tests?
## This tests the litellm support for the openai /generations endpoint
import sys, os
import traceback
from dotenv import load_dotenv
load_dotenv()
import os
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import pytest
import litellm
def test_image_generation_openai():
litellm.set_verbose = True
response = litellm.image_generation(prompt="A cute baby sea otter", model="dall-e-3")
print(f"response: {response}")
# test_image_generation_openai()
# def test_image_generation_azure():
# response = litellm.image_generation(prompt="A cute baby sea otter", api_version="2023-06-01-preview", custom_llm_provider="azure")
# print(f"response: {response}")
# test_image_generation_azure()
# @pytest.mark.asyncio
# async def test_async_image_generation_openai():
# response = litellm.image_generation(prompt="A cute baby sea otter", model="dall-e-3")
# print(f"response: {response}")
# @pytest.mark.asyncio
# async def test_async_image_generation_azure():
# response = litellm.image_generation(prompt="A cute baby sea otter", model="azure/dall-e-3")
# print(f"response: {response}")