Browse Source

Merge pull request #1802 from Yanyutin753/tem-dev

✨feat added environment variables and sync.yml
Timothy Jaeryang Baek 1 year ago
parent
commit
386fc040b1
3 changed files with 20 additions and 6 deletions
  1. 8 4
      backend/apps/images/main.py
  2. 2 1
      backend/apps/rag/main.py
  3. 10 1
      backend/config.py

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

@@ -32,11 +32,15 @@ import logging
 from config import (
     SRC_LOG_LEVELS,
     CACHE_DIR,
+    IMAGE_GENERATION_ENGINE,
     ENABLE_IMAGE_GENERATION,
     AUTOMATIC1111_BASE_URL,
     COMFYUI_BASE_URL,
     IMAGES_OPENAI_API_BASE_URL,
     IMAGES_OPENAI_API_KEY,
+    IMAGE_GENERATION_MODEL,
+    IMAGE_SIZE,
+    IMAGE_STEPS,
 )
 
 
@@ -55,21 +59,21 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-app.state.ENGINE = ""
+app.state.ENGINE = IMAGE_GENERATION_ENGINE
 app.state.ENABLED = ENABLE_IMAGE_GENERATION
 
 app.state.OPENAI_API_BASE_URL = IMAGES_OPENAI_API_BASE_URL
 app.state.OPENAI_API_KEY = IMAGES_OPENAI_API_KEY
 
-app.state.MODEL = ""
+app.state.MODEL = IMAGE_GENERATION_MODEL
 
 
 app.state.AUTOMATIC1111_BASE_URL = AUTOMATIC1111_BASE_URL
 app.state.COMFYUI_BASE_URL = COMFYUI_BASE_URL
 
 
-app.state.IMAGE_SIZE = "512x512"
-app.state.IMAGE_STEPS = 50
+app.state.IMAGE_SIZE = IMAGE_SIZE
+app.state.IMAGE_STEPS = IMAGE_STEPS
 
 
 @app.get("/config")

+ 2 - 1
backend/apps/rag/main.py

@@ -74,6 +74,7 @@ from config import (
     RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE,
     ENABLE_RAG_HYBRID_SEARCH,
     RAG_RERANKING_MODEL,
+    PDF_EXTRACT_IMAGES,
     RAG_RERANKING_MODEL_AUTO_UPDATE,
     RAG_RERANKING_MODEL_TRUST_REMOTE_CODE,
     RAG_OPENAI_API_BASE_URL,
@@ -108,7 +109,7 @@ app.state.RAG_TEMPLATE = RAG_TEMPLATE
 app.state.OPENAI_API_BASE_URL = RAG_OPENAI_API_BASE_URL
 app.state.OPENAI_API_KEY = RAG_OPENAI_API_KEY
 
-app.state.PDF_EXTRACT_IMAGES = False
+app.state.PDF_EXTRACT_IMAGES = PDF_EXTRACT_IMAGES
 
 
 def update_embedding_model(

+ 10 - 1
backend/config.py

@@ -441,6 +441,8 @@ ENABLE_RAG_HYBRID_SEARCH = (
 
 RAG_EMBEDDING_ENGINE = os.environ.get("RAG_EMBEDDING_ENGINE", "")
 
+PDF_EXTRACT_IMAGES = os.environ.get("PDF_EXTRACT_IMAGES", "False").lower() == "true"
+
 RAG_EMBEDDING_MODEL = os.environ.get(
     "RAG_EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2"
 )
@@ -529,18 +531,25 @@ WHISPER_MODEL_AUTO_UPDATE = (
 # Images
 ####################################
 
+IMAGE_GENERATION_ENGINE = os.getenv("IMAGE_GENERATION_ENGINE", "")
+
 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", "")
 
+COMFYUI_BASE_URL = os.getenv("COMFYUI_BASE_URL", "")
 
 IMAGES_OPENAI_API_BASE_URL = os.getenv(
     "IMAGES_OPENAI_API_BASE_URL", OPENAI_API_BASE_URL
 )
 IMAGES_OPENAI_API_KEY = os.getenv("IMAGES_OPENAI_API_KEY", OPENAI_API_KEY)
 
+IMAGE_SIZE = os.getenv("IMAGE_SIZE", "512x512")
+
+IMAGE_STEPS = int(os.getenv("IMAGE_STEPS", 50))
+
+IMAGE_GENERATION_MODEL = os.getenv("IMAGE_GENERATION_MODEL", "")
 
 ####################################
 # Audio