fix upscale api

This commit is contained in:
Vladimir Mandic
2024-02-04 10:14:10 -05:00
parent a93c8160d7
commit 0a0efa3fc1
9 changed files with 147 additions and 54 deletions
+43 -46
View File
@@ -305,10 +305,7 @@ def scrub_dict(dict_obj, keys):
def read_metadata_from_safetensors(filename):
global sd_metadata # pylint: disable=global-statement
if sd_metadata is None:
if not os.path.isfile(sd_metadata_file):
sd_metadata = {}
else:
sd_metadata = shared.readfile(sd_metadata_file, lock=True)
sd_metadata = shared.readfile(sd_metadata_file, lock=True) if os.path.isfile(sd_metadata_file) else {}
res = sd_metadata.get(filename, None)
if res is not None:
return res
@@ -317,48 +314,48 @@ def read_metadata_from_safetensors(filename):
if shared.cmd_opts.no_metadata:
return {}
res = {}
try:
t0 = time.time()
with open(filename, mode="rb") as file:
metadata_len = file.read(8)
metadata_len = int.from_bytes(metadata_len, "little")
json_start = file.read(2)
if metadata_len <= 2 or json_start not in (b'{"', b"{'"):
shared.log.error(f"Not a valid safetensors file: {filename}")
json_data = json_start + file.read(metadata_len-2)
json_obj = json.loads(json_data)
for k, v in json_obj.get("__metadata__", {}).items():
if v.startswith("data:"):
v = 'data'
if k == 'format' and v == 'pt':
continue
large = True if len(v) > 2048 else False
if large and k == 'ss_datasets':
continue
if large and k == 'workflow':
continue
if large and k == 'prompt':
continue
if large and k == 'ss_bucket_info':
continue
if v[0:1] == '{':
try:
v = json.loads(v)
if large and k == 'ss_tag_frequency':
v = { i: len(j) for i, j in v.items() }
if large and k == 'sd_merge_models':
scrub_dict(v, ['sd_merge_recipe'])
except Exception:
pass
res[k] = v
sd_metadata[filename] = res
global sd_metadata_pending # pylint: disable=global-statement
sd_metadata_pending += 1
t1 = time.time()
global sd_metadata_timer # pylint: disable=global-statement
sd_metadata_timer += (t1 - t0)
except Exception as e:
shared.log.error(f"Error reading metadata from: {filename} {e}")
# try:
t0 = time.time()
with open(filename, mode="rb") as file:
metadata_len = file.read(8)
metadata_len = int.from_bytes(metadata_len, "little")
json_start = file.read(2)
if metadata_len <= 2 or json_start not in (b'{"', b"{'"):
shared.log.error(f"Not a valid safetensors file: {filename}")
json_data = json_start + file.read(metadata_len-2)
json_obj = json.loads(json_data)
for k, v in json_obj.get("__metadata__", {}).items():
if v.startswith("data:"):
v = 'data'
if k == 'format' and v == 'pt':
continue
large = True if len(v) > 2048 else False
if large and k == 'ss_datasets':
continue
if large and k == 'workflow':
continue
if large and k == 'prompt':
continue
if large and k == 'ss_bucket_info':
continue
if v[0:1] == '{':
try:
v = json.loads(v)
if large and k == 'ss_tag_frequency':
v = { i: len(j) for i, j in v.items() }
if large and k == 'sd_merge_models':
scrub_dict(v, ['sd_merge_recipe'])
except Exception:
pass
res[k] = v
sd_metadata[filename] = res
global sd_metadata_pending # pylint: disable=global-statement
sd_metadata_pending += 1
t1 = time.time()
global sd_metadata_timer # pylint: disable=global-statement
sd_metadata_timer += (t1 - t0)
# except Exception as e:
# shared.log.error(f"Error reading metadata from: {filename} {e}")
return res