From 3902aa222b00a24f2d7b7158b79efaac9f318923 Mon Sep 17 00:00:00 2001 From: Brendan Hoar Date: Fri, 26 Apr 2024 06:44:41 -0400 Subject: [PATCH 1/6] Better error handling to skip non-standard ss_tag_frequency content --- extensions-builtin/Lora/ui_edit_user_metadata.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/extensions-builtin/Lora/ui_edit_user_metadata.py b/extensions-builtin/Lora/ui_edit_user_metadata.py index 7a07a544e..b6c4d1c6a 100644 --- a/extensions-builtin/Lora/ui_edit_user_metadata.py +++ b/extensions-builtin/Lora/ui_edit_user_metadata.py @@ -21,10 +21,12 @@ re_comma = re.compile(r" *, *") def build_tags(metadata): tags = {} - for _, tags_dict in metadata.get("ss_tag_frequency", {}).items(): - for tag, tag_count in tags_dict.items(): - tag = tag.strip() - tags[tag] = tags.get(tag, 0) + int(tag_count) + ss_tag_frequency = metadata.get("ss_tag_frequency", {}) + if ss_tag_frequency is not None and hasattr(ss_tag_frequency, 'items'): + for _, tags_dict in ss_tag_frequency.items(): + for tag, tag_count in tags_dict.items(): + tag = tag.strip() + tags[tag] = tags.get(tag, 0) + int(tag_count) if tags and is_non_comma_tagset(tags): new_tags = {} From 8dc920228e7c5181cc990845f0febd2ac4b42d87 Mon Sep 17 00:00:00 2001 From: Brendan Hoar Date: Fri, 26 Apr 2024 06:52:21 -0400 Subject: [PATCH 2/6] Better error handling when unable to read metadata from safetensors file --- modules/sd_models.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index ff245b7a6..59742d311 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -280,18 +280,22 @@ def read_metadata_from_safetensors(filename): json_start = file.read(2) assert metadata_len > 2 and json_start in (b'{"', b"{'"), f"{filename} is not a safetensors file" - json_data = json_start + file.read(metadata_len-2) - json_obj = json.loads(json_data) res = {} - for k, v in json_obj.get("__metadata__", {}).items(): - res[k] = v - if isinstance(v, str) and v[0:1] == '{': - try: - res[k] = json.loads(v) - except Exception: - pass + try: + json_data = json_start + file.read(metadata_len-2) + json_obj = json.loads(json_data) + for k, v in json_obj.get("__metadata__", {}).items(): + res[k] = v + if isinstance(v, str) and v[0:1] == '{': + try: + res[k] = json.loads(v) + except Exception: + pass + except: + errors.report(f"Error reading metadata from file: {filename}", exc_info=True) + return res From c5b7559856c5f64792c2425d11890a121497e6bc Mon Sep 17 00:00:00 2001 From: Brendan Hoar Date: Fri, 26 Apr 2024 06:57:32 -0400 Subject: [PATCH 3/6] Better error handling when unable to extract contents of embedding/TI file --- modules/textual_inversion/textual_inversion.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index 253f219c4..dc7833e93 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -181,12 +181,16 @@ class EmbeddingDatabase: else: return - embedding = create_embedding_from_data(data, name, filename=filename, filepath=path) + if data is not None: + embedding = create_embedding_from_data(data, name, filename=filename, filepath=path) - if self.expected_shape == -1 or self.expected_shape == embedding.shape: - self.register_embedding(embedding, shared.sd_model) + if self.expected_shape == -1 or self.expected_shape == embedding.shape: + self.register_embedding(embedding, shared.sd_model) + else: + self.skipped_embeddings[name] = embedding else: - self.skipped_embeddings[name] = embedding + print(f"Unable to load Textual inversion embedding due to data issue: '{name}'.") + def load_from_dir(self, embdir): if not os.path.isdir(embdir.path): From c5ae2254182b803618a4b01c12fa88c42642e806 Mon Sep 17 00:00:00 2001 From: Brendan Hoar Date: Fri, 26 Apr 2024 07:55:39 -0400 Subject: [PATCH 4/6] Better handling of embeddings with two rare, but not unusual, files in them I have encountered pickled embeddings with a short byteorder file at the top-level, as well as a .data/serialization_id file. Both load fine after allowing these files in the dataset. I do not think it is likely adding them to the safe unpickle regular expression would be a security risk, but that's for the maintainers to decide. --- modules/safe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/safe.py b/modules/safe.py index b1d08a792..ee8789268 100644 --- a/modules/safe.py +++ b/modules/safe.py @@ -65,7 +65,7 @@ class RestrictedUnpickler(pickle.Unpickler): # Regular expression that accepts 'dirname/version', 'dirname/data.pkl', and 'dirname/data/' -allowed_zip_names_re = re.compile(r"^([^/]+)/((data/\d+)|version|(data\.pkl))$") +allowed_zip_names_re = re.compile(r"^([^/]+)/((data/\d+)|byteorder|(\.data\/serialization_id)|version|(data\.pkl))$") data_pkl_re = re.compile(r"^([^/]+)/data\.pkl$") def check_zip_filenames(filename, names): From 44afb48447c2ef40f8546fe704bd817881da5a14 Mon Sep 17 00:00:00 2001 From: Brendan Hoar Date: Fri, 26 Apr 2024 08:17:37 -0400 Subject: [PATCH 5/6] Linter fix - extraneous whitespace --- modules/sd_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index 59742d311..06e881207 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -295,7 +295,7 @@ def read_metadata_from_safetensors(filename): pass except: errors.report(f"Error reading metadata from file: {filename}", exc_info=True) - + return res From 60c079995824ebe861029839ee12ca0df6a26e8d Mon Sep 17 00:00:00 2001 From: Brendan Hoar Date: Fri, 26 Apr 2024 08:21:12 -0400 Subject: [PATCH 6/6] Linter - except must not be bare. --- modules/sd_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index 06e881207..06a7cf3f0 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -293,7 +293,7 @@ def read_metadata_from_safetensors(filename): res[k] = json.loads(v) except Exception: pass - except: + except Exception: errors.report(f"Error reading metadata from file: {filename}", exc_info=True) return res