Save model info in json file beside model.
This commit is contained in:
39
__init__.py
39
__init__.py
@@ -60,7 +60,8 @@ preview_extensions = ( # TODO: JavaScript does not know about this (x2 states)
|
|||||||
image_extensions + # order matters
|
image_extensions + # order matters
|
||||||
stable_diffusion_webui_civitai_helper_image_extensions
|
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?
|
#video_extensions = (".avi", ".mp4", ".webm") # TODO: Requires ffmpeg or cv2. Cache preview frame?
|
||||||
|
|
||||||
def split_valid_ext(s, *arg_exts):
|
def split_valid_ext(s, *arg_exts):
|
||||||
@@ -260,6 +261,15 @@ def civitai_get_model_info(sha256_hash):
|
|||||||
return model_response.json()
|
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):
|
def search_web_for_model_url(sha256_hash):
|
||||||
model_info = civitai_get_model_info(sha256_hash)
|
model_info = civitai_get_model_info(sha256_hash)
|
||||||
if len(model_info) > 0:
|
if len(model_info) > 0:
|
||||||
@@ -1046,10 +1056,10 @@ async def get_model_info(request):
|
|||||||
training_comment if training_comment != "None" else ""
|
training_comment if training_comment != "None" else ""
|
||||||
).strip()
|
).strip()
|
||||||
|
|
||||||
info_text_file = abs_name + model_info_extension
|
notes_file = abs_name + model_notes_extension
|
||||||
notes = ""
|
notes = ""
|
||||||
if os.path.isfile(info_text_file):
|
if os.path.isfile(notes_file):
|
||||||
with open(info_text_file, 'r', encoding="utf-8") as f:
|
with open(notes_file, 'r', encoding="utf-8") as f:
|
||||||
notes = f.read()
|
notes = f.read()
|
||||||
|
|
||||||
if metadata is not None:
|
if metadata is not None:
|
||||||
@@ -1136,6 +1146,7 @@ async def download_model(request):
|
|||||||
result["alert"] = "Invalid save path!"
|
result["alert"] = "Invalid save path!"
|
||||||
return web.json_response(result)
|
return web.json_response(result)
|
||||||
|
|
||||||
|
# download model
|
||||||
download_uri = formdata.get("download")
|
download_uri = formdata.get("download")
|
||||||
if download_uri is None:
|
if download_uri is None:
|
||||||
result["alert"] = "Invalid download url!"
|
result["alert"] = "Invalid download url!"
|
||||||
@@ -1159,6 +1170,21 @@ async def download_model(request):
|
|||||||
result["alert"] = "Failed to download model!\n\n" + str(e)
|
result["alert"] = "Failed to download model!\n\n" + str(e)
|
||||||
return web.json_response(result)
|
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")
|
image = formdata.get("image")
|
||||||
if image is not None and image != "":
|
if image is not None and image != "":
|
||||||
try:
|
try:
|
||||||
@@ -1232,7 +1258,7 @@ async def move_model(request):
|
|||||||
return web.json_response(result)
|
return web.json_response(result)
|
||||||
|
|
||||||
# TODO: this could overwrite existing files in destination; do a check beforehand?
|
# 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
|
old_file = old_file_without_extension + extension
|
||||||
if os.path.isfile(old_file):
|
if os.path.isfile(old_file):
|
||||||
new_file = new_file_without_extension + extension
|
new_file = new_file_without_extension + extension
|
||||||
@@ -1286,6 +1312,7 @@ async def delete_model(request):
|
|||||||
print("Deleted file: " + model_path)
|
print("Deleted file: " + model_path)
|
||||||
|
|
||||||
delete_same_name_files(path_and_name, preview_extensions)
|
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,))
|
delete_same_name_files(path_and_name, (model_info_extension,))
|
||||||
|
|
||||||
return web.json_response(result)
|
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_path, model_type = search_path_to_system_path(model_path)
|
||||||
model_extensions = folder_paths_get_supported_pt_extensions(model_type)
|
model_extensions = folder_paths_get_supported_pt_extensions(model_type)
|
||||||
file_path_without_extension, _ = split_valid_ext(model_path, model_extensions)
|
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:
|
if dt_epoch is not None and os.path.exists(filename) and os.path.getmtime(filename) > dt_epoch:
|
||||||
# discard late save
|
# discard late save
|
||||||
|
|||||||
@@ -3333,6 +3333,7 @@ async function getModelInfos(urlText) {
|
|||||||
"fp": file["fp"],
|
"fp": file["fp"],
|
||||||
"quant": file["size"],
|
"quant": file["size"],
|
||||||
"fileFormat": file["format"],
|
"fileFormat": file["format"],
|
||||||
|
"sha256": file["hashes"]["SHA256"],
|
||||||
},
|
},
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
@@ -3353,6 +3354,7 @@ async function getModelInfos(urlText) {
|
|||||||
const infos = hfInfo["modelFiles"].map((file) => {
|
const infos = hfInfo["modelFiles"].map((file) => {
|
||||||
const indexSep = file.lastIndexOf("/");
|
const indexSep = file.lastIndexOf("/");
|
||||||
const filename = file.substring(indexSep + 1);
|
const filename = file.substring(indexSep + 1);
|
||||||
|
// TODO: get sha256 of each HuggingFace model file
|
||||||
return {
|
return {
|
||||||
"images": hfInfo["images"],
|
"images": hfInfo["images"],
|
||||||
"fileName": filename,
|
"fileName": filename,
|
||||||
@@ -3655,6 +3657,7 @@ class DownloadView {
|
|||||||
formData.append("download", info["downloadUrl"]);
|
formData.append("download", info["downloadUrl"]);
|
||||||
formData.append("path", pathDirectory);
|
formData.append("path", pathDirectory);
|
||||||
formData.append("name", modelName);
|
formData.append("name", modelName);
|
||||||
|
formData.append("sha256", info["details"]["sha256"]);
|
||||||
const image = await downloadPreviewSelect.getImage();
|
const image = await downloadPreviewSelect.getImage();
|
||||||
formData.append("image", image === PREVIEW_NONE_URI ? "" : image);
|
formData.append("image", image === PREVIEW_NONE_URI ? "" : image);
|
||||||
formData.append("overwrite", this.elements.overwrite.checked);
|
formData.append("overwrite", this.elements.overwrite.checked);
|
||||||
|
|||||||
Reference in New Issue
Block a user