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"