feat: update available models

src/enums.py

@@ -22,13 +22,8 @@ class ResponseStatusEnum(StrEnum):
 
 
 class ModelEnum(StrEnum):
-    STABLE_DIFFUSION_V1_4 = "stable-diffusion-v1-4"
-    STABLE_DIFFUSION_V1_5 = "stable-diffusion-v1-5"
-    STABLE_DIFFUSION_V2 = "stable-diffusion-v2"
     STABLE_DIFFUSION_V2_1 = "stable-diffusion-v2-1"
-    STABLE_DIFFUSION_V2_768 = "stable-diffusion-v2-768"
     STABLE_DIFFUSION_V2_1_768 = "stable-diffusion-v2-1-768"
-    OPENJOURNEY_V2 = "openjourney-v2"
 
 
 class SchedulerType(StrEnum):

src/event_handlers.py

@@ -7,7 +7,7 @@
 from loguru import logger
 
 from config import celery_settings, firebase_settings
-
+from enums import ModelEnum
 
 def _setup_firebase(app: FastAPI) -> None:
     cred = credentials.Certificate(firebase_settings.firebase_cred_path)
@@ -19,14 +19,11 @@ def _setup_firebase(app: FastAPI) -> None:
 
 def _setup_celery(app: FastAPI) -> None:
     logger.info("Setup Celery")
-    ref = db.reference(f"{firebase_settings.firebase_app_name}/models")
-    data = ref.get()
-    if data is None:
-        exit()
     app.state.celery = {}
-    for model_id in data.keys():
+    for model in ModelEnum:
+        model_id = model.value
         app.state.celery[model_id] = Celery(broker=f"{celery_settings.broker_base_uri}/{model_id}")
-
+        
 
 def start_app_handler(app: FastAPI) -> Callable:
     def startup() -> None: