mirror of
https://github.com/kodxana/madiator-docker-runpod.git
synced 2024-11-25 20:30:11 +01:00
Fix Model Downloader for new 'shared_models' module
This commit is contained in:
parent
f0c15b3315
commit
09b083f24d
2 changed files with 20 additions and 11 deletions
|
@ -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
|
||||||
|
|
|
@ -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"
|
||||||
|
|
Loading…
Reference in a new issue