From 09b083f24df5918d0c5d0f4dc738eaa9f9272c37 Mon Sep 17 00:00:00 2001
From: lutzapps <roman@lutzapps.com>
Date: Sun, 27 Oct 2024 06:30:53 +0700
Subject: [PATCH] Fix Model Downloader for new 'shared_models' module

---
 .../better-ai-launcher/app/app.py             |  7 ++++--
 .../app/utils/model_utils.py                  | 24 ++++++++++++-------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/official-templates/better-ai-launcher/app/app.py b/official-templates/better-ai-launcher/app/app.py
index 0cc72de..eabcf40 100644
--- a/official-templates/better-ai-launcher/app/app.py
+++ b/official-templates/better-ai-launcher/app/app.py
@@ -57,7 +57,7 @@ from utils.shared_models import (
 
 from utils.websocket_utils import send_websocket_message, active_websockets
 from utils.app_configs import get_app_configs, add_app_config, remove_app_config, app_configs
-from utils.model_utils import download_model, check_civitai_url, check_huggingface_url, SHARED_MODELS_DIR, format_size
+from utils.model_utils import download_model, check_civitai_url, check_huggingface_url, format_size #, SHARED_MODELS_DIR # lutzapps - SHARED_MODELS_DIR is owned by shared_models module now
 
 app = Flask(__name__)
 sock = Sock(app)
@@ -628,7 +628,10 @@ def download_model_route():
 @app.route('/get_model_folders')
 def get_model_folders():
     folders = {}
-    for folder in ['Stable-diffusion', 'VAE', 'Lora', 'ESRGAN']:
+    
+    # lutzapps - replace the hard-coded model types
+    for folder, model_type_description in SHARED_MODEL_FOLDERS.items():
+    #for folder in ['Stable-diffusion', 'VAE', 'Lora', 'ESRGAN']:
         folder_path = os.path.join(SHARED_MODELS_DIR, folder)
         if os.path.exists(folder_path):
             total_size = 0
diff --git a/official-templates/better-ai-launcher/app/utils/model_utils.py b/official-templates/better-ai-launcher/app/utils/model_utils.py
index 1c0d3e5..87e168d 100644
--- a/official-templates/better-ai-launcher/app/utils/model_utils.py
+++ b/official-templates/better-ai-launcher/app/utils/model_utils.py
@@ -6,24 +6,30 @@ import json
 import re
 import time
 import math
+# lutzapps - modify for new shared_models module and overwrite for this module
+from shared_models import (ensure_shared_models_folders, SHARED_MODELS_DIR)
 
-SHARED_MODELS_DIR = '/workspace/shared_models'
+#SHARED_MODELS_DIR = '/workspace/shared_models'
 
+# lutzapps - modify this CivitAI model_type mapping to the new SHARED_MODEL_FOLDERS map
 MODEL_TYPE_MAPPING = {
-    'Checkpoint': 'Stable-diffusion',
-    'LORA': 'Lora',
-    'LoCon': 'Lora',
+    # CivitAI-Modeltype: SHARED_MODEL_FOLDERS
+    'Checkpoint': 'ckpt', #'Stable-diffusion', # not clear name for model_type
+    'LORA': 'loras', #'Lora', # now lowercase and plural
+    'LoCon': 'loras', #'Lora', # now lowercase and plural
     'TextualInversion': 'embeddings',
-    'VAE': 'VAE',
+    'VAE': 'vae', #'VAE', # now lowercase
     'Hypernetwork': 'hypernetworks',
-    'AestheticGradient': 'aesthetic_embeddings',
+    'AestheticGradient': 'embeddings', #'aesthetic_embeddings', # store together with "embeddings"
     'ControlNet': 'controlnet',
-    'Upscaler': 'ESRGAN'
+    'Upscaler': 'upscale_models' #'ESRGAN' # there are probably other upscalers not based on ESRGAN
 }
 
 def ensure_shared_folder_exists():
-    for folder in ['Stable-diffusion', 'Lora', 'embeddings', 'VAE', 'hypernetworks', 'aesthetic_embeddings', 'controlnet', 'ESRGAN']:
-        os.makedirs(os.path.join(SHARED_MODELS_DIR, folder), exist_ok=True)
+    # lutzapps - replace with new shared_models code
+    #for folder in ['Stable-diffusion', 'Lora', 'embeddings', 'VAE', 'hypernetworks', 'aesthetic_embeddings', 'controlnet', 'ESRGAN']:
+    #    os.makedirs(os.path.join(SHARED_MODELS_DIR, folder), exist_ok=True)
+    ensure_shared_models_folders()
 
 def check_civitai_url(url):
     prefix = "civitai.com"