mirror of
https://github.com/xtekky/gpt4free.git
synced 2025-12-06 10:40:43 -08:00
77 lines
2.2 KiB
Python
77 lines
2.2 KiB
Python
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:
|
|
# Calling model updates before creating a generator
|
|
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
|