Skip to content

Commit 09f5dd6

Browse files
added the models allowed under this bot
1 parent cecb4f4 commit 09f5dd6

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

utils/models.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
from utils.logger import discord_logger
33

44

5+
ALLOWED_MODELS = {"flux", "klein", "klein-large", "zimage"}
6+
7+
58
async def fetch_and_log_models(config, action: str) -> None:
69
"""
710
Fetch models from API and update config with proper logging.
@@ -14,15 +17,18 @@ async def fetch_and_log_models(config, action: str) -> None:
1417
# Fetch new models from API
1518
models = await initialize_models_async(config)
1619

20+
# Filter to only allowed models
21+
filtered_models = [model for model in models if model in ALLOWED_MODELS]
22+
1723
# Update config with new models (clear and extend to maintain reference)
1824
config.MODELS.clear()
19-
config.MODELS.extend(models)
25+
config.MODELS.extend(filtered_models if filtered_models else [config.image_generation.fallback_model])
2026

2127
# Log successful operation
2228
discord_logger.log_bot_event(
2329
action=action,
2430
status="success",
25-
details=f"Models {action.replace('_', ' ')}: {models}",
31+
details=f"Models {action.replace('_', ' ')}: {filtered_models}",
2632
)
2733

2834
except Exception as e:

0 commit comments

Comments
 (0)