File size: 3,566 Bytes
7a024d2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
# -*- coding: UTF-8 -*-
# handle msg between js and python side
import os
import json
from . import util
from modules import shared
# this is the default root path
root_path = os.getcwd()
# if command line arguement is used to change model folder,
# then model folder is in absolute path, not based on this root path anymore.
# so to make extension work with those absolute model folder paths, model folder also need to be in absolute path
folders = {
"ti": os.path.join(root_path, "embeddings"),
"hyper": os.path.join(root_path, "models", "hypernetworks"),
"ckp": os.path.join(root_path, "models", "Stable-diffusion"),
"lora": os.path.join(root_path, "models", "Lora"),
}
exts = (".bin", ".pt", ".safetensors", ".ckpt")
info_ext = ".info"
vae_suffix = ".vae"
# get cusomter model path
def get_custom_model_folder():
util.printD("Get Custom Model Folder")
global folders
if shared.cmd_opts.embeddings_dir and os.path.isdir(shared.cmd_opts.embeddings_dir):
folders["ti"] = shared.cmd_opts.embeddings_dir
if shared.cmd_opts.hypernetwork_dir and os.path.isdir(shared.cmd_opts.hypernetwork_dir):
folders["hyper"] = shared.cmd_opts.hypernetwork_dir
if shared.cmd_opts.ckpt_dir and os.path.isdir(shared.cmd_opts.ckpt_dir):
folders["ckp"] = shared.cmd_opts.ckpt_dir
if shared.cmd_opts.lora_dir and os.path.isdir(shared.cmd_opts.lora_dir):
folders["lora"] = shared.cmd_opts.lora_dir
# write model info to file
def write_model_info(path, model_info):
util.printD("Write model info to file: " + path)
with open(os.path.realpath(path), 'w') as f:
f.write(json.dumps(model_info, indent=4))
def load_model_info(path):
# util.printD("Load model info from file: " + path)
model_info = None
with open(os.path.realpath(path), 'r') as f:
try:
model_info = json.load(f)
except Exception as e:
util.printD("Selected file is not json: " + path)
util.printD(e)
return
return model_info
# get model file names by model type
# parameter: model_type - string
# return: model name list
def get_model_names_by_type(model_type:str) -> list:
model_folder = folders[model_type]
# get information from filter
# only get those model names don't have a civitai model info file
model_names = []
for root, dirs, files in os.walk(model_folder, followlinks=True):
for filename in files:
item = os.path.join(root, filename)
# check extension
base, ext = os.path.splitext(item)
if ext in exts:
# find a model
model_names.append(filename)
return model_names
# return 2 values: (model_root, model_path)
def get_model_path_by_type_and_name(model_type:str, model_name:str) -> str:
util.printD("Run get_model_path_by_type_and_name")
if model_type not in folders.keys():
util.printD("unknown model_type: " + model_type)
return
if not model_name:
util.printD("model name can not be empty")
return
folder = folders[model_type]
# model could be in subfolder, need to walk.
model_root = ""
model_path = ""
for root, dirs, files in os.walk(folder, followlinks=True):
for filename in files:
if filename == model_name:
# find model
model_root = root
model_path = os.path.join(root, filename)
return (model_root, model_path)
return
|