Madiator2011/madiator-docker-runpod
1
0
Fork 0
mirror of https://github.com/kodxana/madiator-docker-runpod.git synced 2024-11-25 04:10:12 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix Model Downloader for new 'shared_models' module

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

7
official-templates/better-ai-launcher/app/app.py
View file

@ -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

24
official-templates/better-ai-launcher/app/utils/model_utils.py
View file

@ -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"