hoelee/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-08-04 03:10:21 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #1109 from d8ahazard/ModelLoader

Browse Source 
Model Loader, Fixes
This commit is contained in:
AUTOMATIC1111
2022-09-30 09:35:58 +03:00
committed by GitHub
parent f80c3696f6 435fd2112a
commit 25414bcd05
23 changed files with 2176 additions and 1413 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
modules/sd_models.py
View File

@@ -8,7 +8,14 @@ from omegaconf import OmegaConf
from ldm.util import instantiate_from_config
from modules import shared
from modules import shared, modelloader
from modules.paths import models_path
model_dir = "Stable-diffusion"
model_path = os.path.join(models_path, model_dir)
model_name = "sd-v1-4.ckpt"
model_url = "https://drive.yerf.org/wl/?id=EBfTrmcCCUAGaQBXVIj5lJmEhjoP1tgl&mode=grid&download=1"
user_dir = None
CheckpointInfo = namedtuple("CheckpointInfo", ['filename', 'title', 'hash', 'model_name'])
checkpoints_list = {}
@@ -23,14 +30,47 @@ except Exception:
pass
def modeltitle(path, h):
abspath = os.path.abspath(path)
if abspath.startswith(model_dir):
name = abspath.replace(model_dir, '')
else:
name = os.path.basename(path)
if name.startswith("\\") or name.startswith("/"):
name = name[1:]
return f'{name} [{h}]'
def setup_model(dirname):
global model_path
global model_name
global model_url
global user_dir
global model_list
user_dir = dirname
if not os.path.exists(model_path):
os.makedirs(model_path)
checkpoints_list.clear()
list_models()
def checkpoint_tiles():
return sorted([x.title for x in checkpoints_list.values()])
def list_models():
global model_path
global model_url
global model_name
global user_dir
checkpoints_list.clear()
model_dir = os.path.abspath(shared.cmd_opts.ckpt_dir)
model_list = modelloader.load_models(model_path=model_path,model_url=model_url,command_path= user_dir,
ext_filter=[".ckpt"], download_name=model_name)
print(f"Model list: {model_list}")
model_dir = os.path.abspath(model_path)
def modeltitle(path, h):
abspath = os.path.abspath(path)
@@ -53,13 +93,11 @@ def list_models():
title, model_name = modeltitle(cmd_ckpt, h)
checkpoints_list[title] = CheckpointInfo(cmd_ckpt, title, h, model_name)
elif cmd_ckpt is not None and cmd_ckpt != shared.default_sd_model_file:
print(f"Checkpoint in --ckpt argument not found: {cmd_ckpt}", file=sys.stderr)
if os.path.exists(model_dir):
for filename in glob.glob(model_dir + '/**/*.ckpt', recursive=True):
h = model_hash(filename)
title, model_name = modeltitle(filename, h)
checkpoints_list[title] = CheckpointInfo(filename, title, h, model_name)
print(f"Checkpoint in --ckpt argument not found (Possible it was moved to {model_path}: {cmd_ckpt}", file=sys.stderr)
for filename in model_list:
h = model_hash(filename)
title = modeltitle(filename, h)
checkpoints_list[title] = CheckpointInfo(filename, title, h, model_name)
def get_closet_checkpoint_match(searchString):
applicable = sorted([info for info in checkpoints_list.values() if searchString in info.title], key = lambda x:len(x.title))
@@ -69,6 +107,7 @@ def get_closet_checkpoint_match(searchString):
def model_hash(filename):
try:
print(f"Opening: {filename}")
with open(filename, "rb") as file:
import hashlib
m = hashlib.sha256()
@@ -89,7 +128,7 @@ def select_checkpoint():
if len(checkpoints_list) == 0:
print(f"No checkpoints found. When searching for checkpoints, looked at:", file=sys.stderr)
print(f" - file {os.path.abspath(shared.cmd_opts.ckpt)}", file=sys.stderr)
print(f" - directory {os.path.abspath(shared.cmd_opts.ckpt_dir)}", file=sys.stderr)
print(f" - directory {os.path.abspath(shared.cmd_opts.stablediffusion_models_path)}", file=sys.stderr)
print(f"Can't run without a checkpoint. Find and place a .ckpt file into any of those locations. The program will exit.", file=sys.stderr)
exit(1)