holm/easydiffusion
1
0
Fork 0
forked from extern/easydiffusion
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow named models in the dropdown

Browse Source
This commit is contained in:
cmdr2 2023-08-01 16:16:38 +05:30
parent ee6db85768
commit eba7bab15e
2 changed files with 23 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
ui/easydiffusion/model_manager.py
View File

@ -290,11 +290,11 @@ def is_malicious_model(file_path):
def getModels(scan_for_malicious: bool = True):
models = {
"options": {
"stable-diffusion": ["sd-v1-4"],
"stable-diffusion": [{"sd-v1-4": "SD 1.4"}],
"vae": [],
"hypernetwork": [],
"lora": [],
"codeformer": ["codeformer"],
"codeformer": [{"codeformer": "CodeFormer"}],
"embeddings": [],
"controlnet": [],
},
@ -305,9 +305,9 @@ def getModels(scan_for_malicious: bool = True):
class MaliciousModelException(Exception):
"Raised when picklescan reports a problem with a model"
def scan_directory(directory, suffixes, directoriesFirst: bool = True):
def scan_directory(directory, suffixes, directoriesFirst: bool = True, default_entries=[]):
tree = list(default_entries)
nonlocal models_scanned
tree = []
for entry in sorted(
os.scandir(directory),
key=lambda entry: (entry.is_file() == directoriesFirst, entry.name.lower()),
@ -326,7 +326,14 @@ def getModels(scan_for_malicious: bool = True):
raise MaliciousModelException(entry.path)
if scan_for_malicious:
known_models[entry.path] = mtime
tree.append(entry.name[: -len(matching_suffix)])
model_id = entry.name[: -len(matching_suffix)]
model_exists = False
for m in tree: # allows default "named" models, like CodeFormer and known ControlNet models
if (isinstance(m, str) and model_id == m) or (isinstance(m, dict) and model_id in m):
model_exists = True
break
if not model_exists:
tree.append(model_id)
elif entry.is_dir():
scan = scan_directory(entry.path, suffixes, directoriesFirst=False)
@ -343,7 +350,8 @@ def getModels(scan_for_malicious: bool = True):
os.makedirs(models_dir)
try:
models["options"][model_type] = scan_directory(models_dir, model_extensions)
default_tree = models["options"].get(model_type, [])
models["options"][model_type] = scan_directory(models_dir, model_extensions, default_entries=default_tree)
except MaliciousModelException as e:
models["scan-error"] = str(e)

28
ui/media/js/searchable-models.js
View File

@ -552,17 +552,23 @@ class ModelDropdown {
this.createModelNodeList(`${folderName || ""}/${childFolderName}`, childModels, false)
)
} else {
let modelId = model
let modelName = model
if (typeof model === "object") {
modelId = Object.keys(model)[0]
modelName = model[modelId]
}
const classes = ["model-file"]
if (isRootFolder) {
classes.push("in-root-folder")
}
// Remove the leading slash from the model path
const fullPath = folderName ? `${folderName.substring(1)}/${model}` : model
const fullPath = folderName ? `${folderName.substring(1)}/${modelId}` : modelId
modelsMap.set(
model,
modelId,
createElement("li", { "data-path": fullPath }, classes, [
createElement("i", undefined, ["fa-regular", "fa-file", "icon"]),
model,
modelName,
])
)
}
@ -643,22 +649,6 @@ async function getModels(scanForMalicious = true) {
makeImageBtn.disabled = true
}
/* This code should no longer be needed. Commenting out for now, will cleanup later.
const sd_model_setting_key = "stable_diffusion_model"
const vae_model_setting_key = "vae_model"
const hypernetwork_model_key = "hypernetwork_model"
const stableDiffusionOptions = modelsOptions['stable-diffusion']
const vaeOptions = modelsOptions['vae']
const hypernetworkOptions = modelsOptions['hypernetwork']
// TODO: set default for model here too
SETTINGS[sd_model_setting_key].default = stableDiffusionOptions[0]
if (getSetting(sd_model_setting_key) == '' || SETTINGS[sd_model_setting_key].value == '') {
setSetting(sd_model_setting_key, stableDiffusionOptions[0])
}
*/
// notify ModelDropdown objects to refresh
document.dispatchEvent(new Event("refreshModels"))
} catch (e) {