Skip to content

Commit

Permalink
refactor(g4f/Provider/PollinationsImage.py): Improve model synchroniz…
Browse files Browse the repository at this point in the history
…ation and initialization
  • Loading branch information
kqlio67 committed Feb 21, 2025
1 parent f2e3f2e commit be58926
Showing 1 changed file with 15 additions and 3 deletions.
18 changes: 15 additions & 3 deletions g4f/Provider/PollinationsImage.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,22 @@ class PollinationsImage(PollinationsAI):
default_model = "flux"
default_vision_model = None
default_image_model = default_model
image_models = []
image_models = [default_image_model] # Default models
_models_loaded = False # Add a checkbox for synchronization

@classmethod
def get_models(cls, **kwargs):
if not cls.image_models:
cls.image_models = list(dict.fromkeys([*cls.image_models, *cls.extra_image_models]))
if not cls._models_loaded:
# Calling the parent method to load models
super().get_models()
# Combine models from the parent class and additional ones
all_image_models = list(dict.fromkeys(
cls.image_models +
PollinationsAI.image_models +
cls.extra_image_models
))
cls.image_models = all_image_models
cls._models_loaded = True
return cls.image_models

@classmethod
Expand All @@ -36,6 +46,8 @@ async def create_async_generator(
safe: bool = False,
**kwargs
) -> AsyncResult:
# Calling model updates before creating a generator
cls.get_models()
async for chunk in cls._generate_image(
model=model,
prompt=format_image_prompt(messages, prompt),
Expand Down

0 comments on commit be58926

Please # to comment.