Skip to content

Commit

Permalink
Display name of error latent file
Browse files Browse the repository at this point in the history
When trying to load stored latents, if an error occurs, this change will tell you what file failed to load
Currently it will just tell you that something failed without telling you which file
  • Loading branch information
Cauldrath committed Apr 19, 2024
1 parent 71e2c91 commit feefcf2
Showing 1 changed file with 12 additions and 9 deletions.
21 changes: 12 additions & 9 deletions library/train_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -2123,18 +2123,21 @@ def is_disk_cached_latents_is_expected(reso, npz_path: str, flip_aug: bool):
if not os.path.exists(npz_path):
return False

npz = np.load(npz_path)
if "latents" not in npz or "original_size" not in npz or "crop_ltrb" not in npz: # old ver?
return False
if npz["latents"].shape[1:3] != expected_latents_size:
return False

if flip_aug:
if "latents_flipped" not in npz:
try:
npz = np.load(npz_path)
if "latents" not in npz or "original_size" not in npz or "crop_ltrb" not in npz: # old ver?
return False
if npz["latents_flipped"].shape[1:3] != expected_latents_size:
if npz["latents"].shape[1:3] != expected_latents_size:
return False

if flip_aug:
if "latents_flipped" not in npz:
return False
if npz["latents_flipped"].shape[1:3] != expected_latents_size:
return False
except:
raise RuntimeError(f"Error loading file: {npz_path}")

return True


Expand Down

0 comments on commit feefcf2

Please sign in to comment.