added unified layer for the image based models

Author: Circuit-Overtime

config.py

@@ -12,6 +12,9 @@
 
 load_dotenv(override=True)
 
+# Allowed image generation models
+ALLOWED_MODELS = {"flux", "klein", "klein-large", "zimage"}
+
 
 class BotConfig(BaseModel):
     command_prefix: str

utils/image_gen_utils.py

@@ -7,15 +7,13 @@
 import json
 from PIL import Image
 from exceptions import PromptTooLongError, DimensionTooSmallError, APIError
-from config import config
+from config import config, ALLOWED_MODELS
 import os 
 from dotenv import load_dotenv
 load_dotenv()
 
 __all__: list[str] = ("generate_image", "validate_prompt", "validate_dimensions", "validate_model")
 
-ALLOWED_MODELS = {"flux", "klein", "klein-large", "zimage"}
-
 
 def validate_prompt(prompt) -> None:
     if len(prompt) > config.image_generation.validation.max_prompt_length:

utils/models.py

@@ -1,10 +1,7 @@
-from config import initialize_models_async
+from config import initialize_models_async, ALLOWED_MODELS
 from utils.logger import discord_logger
 
 
-ALLOWED_MODELS = {"flux", "klein", "klein-large", "zimage"}
-
-
 async def fetch_and_log_models(config, action: str) -> None:
     """
     Fetch models from API and update config with proper logging.