-
-
Notifications
You must be signed in to change notification settings - Fork 13.4k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add OpenRouter and DeepInfraImage Provider (#1814)
- Loading branch information
Showing
14 changed files
with
164 additions
and
38 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,74 @@ | ||
from __future__ import annotations | ||
|
||
import requests | ||
|
||
from .base_provider import AsyncGeneratorProvider, ProviderModelMixin | ||
from ..typing import AsyncResult, Messages | ||
from ..requests import StreamSession, raise_for_status | ||
from ..image import ImageResponse | ||
|
||
class DeepInfraImage(AsyncGeneratorProvider, ProviderModelMixin): | ||
url = "https://deepinfra.com" | ||
working = True | ||
default_model = 'stability-ai/sdxl' | ||
|
||
@classmethod | ||
def get_models(cls): | ||
if not cls.models: | ||
url = 'https://api.deepinfra.com/models/featured' | ||
models = requests.get(url).json() | ||
cls.models = [model['model_name'] for model in models if model["reported_type"] == "text-to-image"] | ||
return cls.models | ||
|
||
@classmethod | ||
async def create_async_generator( | ||
cls, | ||
model: str, | ||
messages: Messages, | ||
**kwargs | ||
) -> AsyncResult: | ||
yield await cls.create_async(messages[-1]["content"], model, **kwargs) | ||
|
||
@classmethod | ||
async def create_async( | ||
cls, | ||
prompt: str, | ||
model: str, | ||
api_key: str = None, | ||
api_base: str = "https://api.deepinfra.com/v1/inference", | ||
proxy: str = None, | ||
timeout: int = 180, | ||
extra_data: dict = {}, | ||
**kwargs | ||
) -> ImageResponse: | ||
headers = { | ||
'Accept-Encoding': 'gzip, deflate, br', | ||
'Accept-Language': 'en-US', | ||
'Connection': 'keep-alive', | ||
'Origin': 'https://deepinfra.com', | ||
'Referer': 'https://deepinfra.com/', | ||
'Sec-Fetch-Dest': 'empty', | ||
'Sec-Fetch-Mode': 'cors', | ||
'Sec-Fetch-Site': 'same-site', | ||
'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36', | ||
'X-Deepinfra-Source': 'web-embed', | ||
'sec-ch-ua': '"Google Chrome";v="119", "Chromium";v="119", "Not?A_Brand";v="24"', | ||
'sec-ch-ua-mobile': '?0', | ||
'sec-ch-ua-platform': '"macOS"', | ||
} | ||
if api_key is not None: | ||
headers["Authorization"] = f"Bearer {api_key}" | ||
async with StreamSession( | ||
proxies={"all": proxy}, | ||
headers=headers, | ||
timeout=timeout | ||
) as session: | ||
model = cls.get_model(model) | ||
data = {"prompt": prompt, **extra_data} | ||
data = {"input": data} if model == cls.default_model else data | ||
async with session.post(f"{api_base.rstrip('/')}/{model}", json=data) as response: | ||
await raise_for_status(response) | ||
data = await response.json() | ||
images = data["output"] if "output" in data else data["images"] | ||
images = images[0] if len(images) == 1 else images | ||
return ImageResponse(images, prompt) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,31 @@ | ||
from __future__ import annotations | ||
|
||
import requests | ||
|
||
from .Openai import Openai | ||
from ...typing import AsyncResult, Messages | ||
|
||
class OpenRouter(Openai): | ||
url = "https://openrouter.ai" | ||
working = True | ||
default_model = "openrouter/auto" | ||
|
||
@classmethod | ||
def get_models(cls): | ||
if not cls.models: | ||
url = 'https://openrouter.ai/api/v1/models' | ||
models = requests.get(url).json()["data"] | ||
cls.models = [model['id'] for model in models] | ||
return cls.models | ||
|
||
@classmethod | ||
def create_async_generator( | ||
cls, | ||
model: str, | ||
messages: Messages, | ||
api_base: str = "https://openrouter.ai/api/v1", | ||
**kwargs | ||
) -> AsyncResult: | ||
return super().create_async_generator( | ||
model, messages, api_base=api_base, **kwargs | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.