TBNilles/ComfyUI-Model-Manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added "Download Model Info" button to model info view.

Browse Source 
- Always rehashes model (slow for big files).
- Refactored model info download into function on server.
This commit is contained in:
Christian Bastian
2024-09-22 21:13:54 -04:00
parent 1ddba41cc0
commit 3b8735afef
2 changed files with 97 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
__init__.py
View File

@@ -963,7 +963,48 @@ async def get_directory_list(request):
return web.json_response(dir_list)
@server.PromptServer.instance.routes.post("/model-manager/models/scan-download")
def try_download_and_save_model_info(model_file_path):
success = (0, 0, 0) #info, notes, url
head, _ = os.path.splitext(model_file_path)
model_info_path = head + model_info_extension
model_notes_path = head + model_notes_extension
model_url_path = head + ".url"
if os.path.exists(model_info_path) and os.path.exists(model_notes_path) and os.path.exists(model_url_path):
return success
print("Scanning " + model_file_path)
model_info = {}
model_info = ModelInfo.search_info(model_file_path, cache=True, use_cached=True)
if len(model_info) == 0:
return success
success[0] = 1
if not os.path.exists(model_notes_path):
notes = ModelInfo.search_notes(model_file_path)
if not notes.isspace() and notes != "":
try:
with open(model_notes_path, "w", encoding="utf-8") as f:
f.write(notes)
print("Saved file: " + model_notes_path)
success[1] = 1
except Exception as e:
print(f"Failed to save {model_notes_path}!")
print(e, file=sys.stderr, flush=True)
if not os.path.exists(model_url_path):
web_url = ModelInfo.get_url(model_info)
if web_url is not None and web_url != "":
try:
save_web_url(model_url_path, web_url)
print("Saved file: " + model_url_path)
success[2] = 1
except Exception as e:
print(f"Failed to save {model_url_path}!")
print(e, file=sys.stderr, flush=True)
return success
@server.PromptServer.instance.routes.post("/model-manager/models/scan")
async def try_scan_download(request):
refresh = request.query.get("refresh", None) is not None
response = {
@@ -981,42 +1022,10 @@ async def try_scan_download(request):
if file_extension not in model_extension_whitelist:
continue
model_file_path = root + os.path.sep + file
model_info_path = root + os.path.sep + file_name + model_info_extension
model_notes_path = root + os.path.sep + file_name + model_notes_extension
model_url_path = root + os.path.sep + file_name + ".url"
if os.path.exists(model_info_path) and os.path.exists(model_notes_path) and os.path.exists(model_url_path):
continue
print("Scanning " + model_file_path)
model_info = {}
model_info = ModelInfo.search_info(model_file_path, cache=True, use_cached=True)
if len(model_info) == 0:
continue
response["infoCount"] += 1
if not os.path.exists(model_notes_path):
notes = ModelInfo.search_notes(model_file_path)
if not notes.isspace() and notes != "":
try:
with open(model_notes_path, "w", encoding="utf-8") as f:
f.write(notes)
print("Saved file: " + model_notes_path)
response["notesCount"] += 1
except Exception as e:
print(f"Failed to save {model_notes_path}!")
print(e, file=sys.stderr, flush=True)
if not os.path.exists(model_url_path):
web_url = ModelInfo.get_url(model_info)
if web_url is not None and web_url != "":
try:
save_web_url(model_url_path, web_url)
print("Saved file: " + model_url_path)
response["urlCount"] += 1
except Exception as e:
print(f"Failed to save {model_url_path}!")
print(e, file=sys.stderr, flush=True)
savedInfo, savedNotes, savedUrl = try_download_and_save_model_info(model_file_path)
response["infoCount"] += savedInfo
response["notesCount"] += savedNotes
response["urlCount"] += savedUrl
response["success"] = True
return web.json_response(response)
@@ -1295,6 +1304,28 @@ async def get_system_separator(request):
return web.json_response(os.path.sep)
@server.PromptServer.instance.routes.post("/model-manager/model/download/info")
async def download_model_info(request):
result = { "success": False }
model_path = request.query.get("path", None)
if model_path is None:
result["alert"] = "Missing model path!"
return web.json_response(result)
model_path = urllib.parse.unquote(model_path)
abs_path, model_type = search_path_to_system_path(model_path)
if abs_path is None:
result["alert"] = "Invalid model path!"
return web.json_response(result)
model_info = ModelInfo.search_info(abs_path, cache=True, use_cached=False)
if len(model_info) > 0:
result["success"] = True
return web.json_response(result)
@server.PromptServer.instance.routes.post("/model-manager/model/download")
async def download_model(request):
formdata = await request.post()

31
web/model-manager.js
View File

@@ -2869,6 +2869,33 @@ class ModelInfo {
button.disabled = false;
},
}).element,
new ComfyButton({
icon: 'earth-arrow-down',
tooltip: 'Try download model info.',
classList: 'comfyui-button icon-button',
action: async(e) => {
const [button, icon, span] = comfyButtonDisambiguate(e.target);
button.disabled = true;
const success = await comfyRequest(
`/model-manager/model/download/info?path=${path}`,
{
method: 'POST',
body: {},
}
).then((data) => {
const success = data['success'];
const message = data['alert'];
if (message !== undefined) {
window.alert(message);
}
return success;
}).catch((err) => {
return false;
});
comfyButtonAlert(e.target, success, 'mdi-check-bold', 'mdi-close-thick');
button.disabled = false;
},
}).element,
]),
$el('div.row.tab-header', [
$el('div.row.tab-header-flex-block', [
@@ -4722,7 +4749,7 @@ class SettingsView {
const [button, icon, span] = comfyButtonDisambiguate(e.target);
button.disabled = true;
const data = await comfyRequest('/model-manager/models/scan-download', {
const data = await comfyRequest('/model-manager/models/scan', {
method: 'POST',
body: JSON.stringify({}),
}).catch((err) => {
@@ -4732,7 +4759,7 @@ class SettingsView {
const infoCount = data['infoCount'];
const notesCount = data['notesCount'];
const urlCount = data['urlCount'];
window.alert(`${successMessage}\n\nInfo Count: ${infoCount}\nNotes Count: ${notesCount}\nUrl Count: ${urlCount}`);
window.alert(`${successMessage}\nScanned: ${infoCount}\nSaved Notes: ${notesCount}\nSaved Url: ${urlCount}`);
comfyButtonAlert(e.target, success);
if (infoCount > 0 || notesCount > 0 || urlCount > 0) {
await this.reload(true);