Save model info in json file beside model.

This commit is contained in:
Christian Bastian
2024-09-11 17:34:29 -04:00
parent 35878d8514
commit f15c5f9071
2 changed files with 36 additions and 6 deletions

View File

@@ -60,7 +60,8 @@ preview_extensions = ( # TODO: JavaScript does not know about this (x2 states)
image_extensions + # order matters
stable_diffusion_webui_civitai_helper_image_extensions
)
model_info_extension = ".txt"
model_notes_extension = ".txt"
model_info_extension = ".json"
#video_extensions = (".avi", ".mp4", ".webm") # TODO: Requires ffmpeg or cv2. Cache preview frame?
def split_valid_ext(s, *arg_exts):
@@ -260,6 +261,15 @@ def civitai_get_model_info(sha256_hash):
return model_response.json()
def search_web_for_model_info(sha256_hash):
model_info = civitai_get_model_info(sha256_hash)
if len(model_info) > 0: return model_info
# TODO: search other websites
return {}
def search_web_for_model_url(sha256_hash):
model_info = civitai_get_model_info(sha256_hash)
if len(model_info) > 0:
@@ -1046,10 +1056,10 @@ async def get_model_info(request):
training_comment if training_comment != "None" else ""
).strip()
info_text_file = abs_name + model_info_extension
notes_file = abs_name + model_notes_extension
notes = ""
if os.path.isfile(info_text_file):
with open(info_text_file, 'r', encoding="utf-8") as f:
if os.path.isfile(notes_file):
with open(notes_file, 'r', encoding="utf-8") as f:
notes = f.read()
if metadata is not None:
@@ -1136,6 +1146,7 @@ async def download_model(request):
result["alert"] = "Invalid save path!"
return web.json_response(result)
# download model
download_uri = formdata.get("download")
if download_uri is None:
result["alert"] = "Invalid download url!"
@@ -1159,6 +1170,21 @@ async def download_model(request):
result["alert"] = "Failed to download model!\n\n" + str(e)
return web.json_response(result)
# download model info
sha256_hash = formdata.get("sha256", None)
if sha256_hash is not None:
model_info = search_web_for_model_info(sha256_hash)
if len(model_info) > 0:
info_path = os.path.splitext(file_name)[0] + ".json"
try:
with open(info_path, "w", encoding="utf-8") as f:
json.dump(model_info, f, indent=4)
print("Saved file: " + info_path)
except ValueError as e:
print(e, file=sys.stderr, flush=True)
result["alert"] = "Failed to save model info!\n\n" + str(e) # TODO: >1 alert? concat?
# save image as model preview
image = formdata.get("image")
if image is not None and image != "":
try:
@@ -1232,7 +1258,7 @@ async def move_model(request):
return web.json_response(result)
# TODO: this could overwrite existing files in destination; do a check beforehand?
for extension in preview_extensions + (model_info_extension,):
for extension in preview_extensions + (model_notes_extension,) + (model_info_extension,):
old_file = old_file_without_extension + extension
if os.path.isfile(old_file):
new_file = new_file_without_extension + extension
@@ -1286,6 +1312,7 @@ async def delete_model(request):
print("Deleted file: " + model_path)
delete_same_name_files(path_and_name, preview_extensions)
delete_same_name_files(path_and_name, (model_notes_extension,))
delete_same_name_files(path_and_name, (model_info_extension,))
return web.json_response(result)
@@ -1310,7 +1337,7 @@ async def set_notes(request):
model_path, model_type = search_path_to_system_path(model_path)
model_extensions = folder_paths_get_supported_pt_extensions(model_type)
file_path_without_extension, _ = split_valid_ext(model_path, model_extensions)
filename = os.path.normpath(file_path_without_extension + model_info_extension)
filename = os.path.normpath(file_path_without_extension + model_notes_extension)
if dt_epoch is not None and os.path.exists(filename) and os.path.getmtime(filename) > dt_epoch:
# discard late save

View File

@@ -3333,6 +3333,7 @@ async function getModelInfos(urlText) {
"fp": file["fp"],
"quant": file["size"],
"fileFormat": file["format"],
"sha256": file["hashes"]["SHA256"],
},
});
});
@@ -3353,6 +3354,7 @@ async function getModelInfos(urlText) {
const infos = hfInfo["modelFiles"].map((file) => {
const indexSep = file.lastIndexOf("/");
const filename = file.substring(indexSep + 1);
// TODO: get sha256 of each HuggingFace model file
return {
"images": hfInfo["images"],
"fileName": filename,
@@ -3655,6 +3657,7 @@ class DownloadView {
formData.append("download", info["downloadUrl"]);
formData.append("path", pathDirectory);
formData.append("name", modelName);
formData.append("sha256", info["details"]["sha256"]);
const image = await downloadPreviewSelect.getImage();
formData.append("image", image === PREVIEW_NONE_URI ? "" : image);
formData.append("overwrite", this.elements.overwrite.checked);