Bläddra i källkod

Merge pull request #1500 from lainedfles/config-image-generation-enabled

feat: add IMAGE_GENERATION_ENABLED env var
Timothy Jaeryang Baek 1 år sedan
förälder
incheckning
916ce9bc4f
2 ändrade filer med 11 tillägg och 2 borttagningar
  1. 8 2
      backend/apps/images/main.py
  2. 3 0
      backend/config.py

+ 8 - 2
backend/apps/images/main.py

@@ -29,7 +29,13 @@ import base64
 import json
 import logging
 
-from config import SRC_LOG_LEVELS, CACHE_DIR, AUTOMATIC1111_BASE_URL, COMFYUI_BASE_URL
+from config import (
+    SRC_LOG_LEVELS,
+    CACHE_DIR,
+    IMAGE_GENERATION_ENABLED,
+    AUTOMATIC1111_BASE_URL,
+    COMFYUI_BASE_URL,
+)
 
 
 log = logging.getLogger(__name__)
@@ -48,7 +54,7 @@ app.add_middleware(
 )
 
 app.state.ENGINE = ""
-app.state.ENABLED = False
+app.state.ENABLED = IMAGE_GENERATION_ENABLED
 
 app.state.OPENAI_API_KEY = ""
 app.state.MODEL = ""

+ 3 - 0
backend/config.py

@@ -456,5 +456,8 @@ WHISPER_MODEL_DIR = os.getenv("WHISPER_MODEL_DIR", f"{CACHE_DIR}/whisper/models"
 # Images
 ####################################
 
+ENABLE_IMAGE_GENERATION = (
+    os.environ.get("ENABLE_IMAGE_GENERATION", "").lower() == "true"
+)
 AUTOMATIC1111_BASE_URL = os.getenv("AUTOMATIC1111_BASE_URL", "")
 COMFYUI_BASE_URL = os.getenv("COMFYUI_BASE_URL", "")