|
|
from __future__ import annotations |
|
|
|
|
|
from typing import Optional |
|
|
|
|
|
from .helper import format_media_prompt |
|
|
from ..typing import AsyncResult, Messages, MediaListType |
|
|
from ..config import STATIC_URL |
|
|
from .PollinationsAI import PollinationsAI |
|
|
|
|
|
class PollinationsImage(PollinationsAI): |
|
|
label = "PollinationsImage" |
|
|
parent = PollinationsAI.__name__ |
|
|
active_by_default = False |
|
|
default_model = "flux" |
|
|
default_vision_model = None |
|
|
default_image_model = default_model |
|
|
audio_models = {} |
|
|
|
|
|
@classmethod |
|
|
def get_models(cls, **kwargs): |
|
|
PollinationsAI.get_models() |
|
|
cls.image_models = PollinationsAI.image_models |
|
|
cls.models = cls.image_models |
|
|
return cls.models |
|
|
|
|
|
@classmethod |
|
|
def get_grouped_models(cls) -> dict[str, list[str]]: |
|
|
PollinationsAI.get_models() |
|
|
return [ |
|
|
{"group": "Image Generation", "models": PollinationsAI.image_models}, |
|
|
] |
|
|
|
|
|
@classmethod |
|
|
async def create_async_generator( |
|
|
cls, |
|
|
model: str, |
|
|
messages: Messages, |
|
|
media: MediaListType = None, |
|
|
proxy: str = None, |
|
|
referrer: str = STATIC_URL, |
|
|
api_key: str = None, |
|
|
prompt: str = None, |
|
|
aspect_ratio: str = None, |
|
|
width: int = None, |
|
|
height: int = None, |
|
|
seed: Optional[int] = None, |
|
|
cache: bool = False, |
|
|
nologo: bool = True, |
|
|
private: bool = False, |
|
|
enhance: bool = False, |
|
|
safe: bool = False, |
|
|
transparent: bool = False, |
|
|
n: int = 1, |
|
|
**kwargs |
|
|
) -> AsyncResult: |
|
|
|
|
|
cls.get_models() |
|
|
async for chunk in cls._generate_image( |
|
|
model=model, |
|
|
prompt=format_media_prompt(messages, prompt), |
|
|
media=media, |
|
|
proxy=proxy, |
|
|
aspect_ratio=aspect_ratio, |
|
|
width=width, |
|
|
height=height, |
|
|
seed=seed, |
|
|
cache=cache, |
|
|
nologo=nologo, |
|
|
private=private, |
|
|
enhance=enhance, |
|
|
safe=safe, |
|
|
transparent=transparent, |
|
|
n=n, |
|
|
referrer=referrer, |
|
|
api_key=api_key |
|
|
): |
|
|
yield chunk |
|
|
|
