Fix Model Downloader for new 'shared_models' module

This commit is contained in:
lutzapps 2024-10-27 06:30:53 +07:00
parent f0c15b3315
commit 09b083f24d
2 changed files with 20 additions and 11 deletions

View file

@ -57,7 +57,7 @@ from utils.shared_models import (
from utils.websocket_utils import send_websocket_message, active_websockets 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.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__) app = Flask(__name__)
sock = Sock(app) sock = Sock(app)
@ -628,7 +628,10 @@ def download_model_route():
@app.route('/get_model_folders') @app.route('/get_model_folders')
def get_model_folders(): def get_model_folders():
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) folder_path = os.path.join(SHARED_MODELS_DIR, folder)
if os.path.exists(folder_path): if os.path.exists(folder_path):
total_size = 0 total_size = 0

View file

@ -6,24 +6,30 @@ import json
import re import re
import time import time
import math 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 = { MODEL_TYPE_MAPPING = {
'Checkpoint': 'Stable-diffusion', # CivitAI-Modeltype: SHARED_MODEL_FOLDERS
'LORA': 'Lora', 'Checkpoint': 'ckpt', #'Stable-diffusion', # not clear name for model_type
'LoCon': 'Lora', 'LORA': 'loras', #'Lora', # now lowercase and plural
'LoCon': 'loras', #'Lora', # now lowercase and plural
'TextualInversion': 'embeddings', 'TextualInversion': 'embeddings',
'VAE': 'VAE', 'VAE': 'vae', #'VAE', # now lowercase
'Hypernetwork': 'hypernetworks', 'Hypernetwork': 'hypernetworks',
'AestheticGradient': 'aesthetic_embeddings', 'AestheticGradient': 'embeddings', #'aesthetic_embeddings', # store together with "embeddings"
'ControlNet': 'controlnet', 'ControlNet': 'controlnet',
'Upscaler': 'ESRGAN' 'Upscaler': 'upscale_models' #'ESRGAN' # there are probably other upscalers not based on ESRGAN
} }
def ensure_shared_folder_exists(): def ensure_shared_folder_exists():
for folder in ['Stable-diffusion', 'Lora', 'embeddings', 'VAE', 'hypernetworks', 'aesthetic_embeddings', 'controlnet', 'ESRGAN']: # lutzapps - replace with new shared_models code
os.makedirs(os.path.join(SHARED_MODELS_DIR, folder), exist_ok=True) #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): def check_civitai_url(url):
prefix = "civitai.com" prefix = "civitai.com"