From 3b8ea55bc8ec0cd00b31001e431021c39b174150 Mon Sep 17 00:00:00 2001
From: Self Denial <selfdenial@pm.me>
Date: Wed, 10 Apr 2024 23:21:12 -0600
Subject: [PATCH] Add IMAGE_GENERATION_ENABLED env var

With default of `False`.
---
 backend/apps/images/main.py | 4 ++--
 backend/config.py           | 3 +++
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/backend/apps/images/main.py b/backend/apps/images/main.py
index af8cbf7c5..bf3125985 100644
--- a/backend/apps/images/main.py
+++ b/backend/apps/images/main.py
@@ -29,7 +29,7 @@ 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 +48,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 = ""
diff --git a/backend/config.py b/backend/config.py
index 6e3cf92a9..24ccaa479 100644
--- a/backend/config.py
+++ b/backend/config.py
@@ -452,5 +452,8 @@ WHISPER_MODEL_DIR = os.getenv("WHISPER_MODEL_DIR", f"{CACHE_DIR}/whisper/models"
 # Images
 ####################################
 
+IMAGE_GENERATION_ENABLED = (
+    os.environ.get("IMAGE_GENERATION_ENABLED", "").lower() == "true"
+)
 AUTOMATIC1111_BASE_URL = os.getenv("AUTOMATIC1111_BASE_URL", "")
 COMFYUI_BASE_URL = os.getenv("COMFYUI_BASE_URL", "")