generalise to popular lossless formats

This commit is contained in:
DepFA 2022-10-14 18:23:20 +01:00 committed by AUTOMATIC1111
parent 9a1dcd78ed
commit ddf6899df0
1 changed files with 2 additions and 6 deletions

View File

@ -88,18 +88,14 @@ class EmbeddingDatabase:
data = [] data = []
if filename.upper().endswith('.PNG'): if os.path.splitext(filename.upper())[-1] in ['.PNG', '.WEBP', '.JXL', '.AVIF']:
embed_image = Image.open(path) embed_image = Image.open(path)
if 'sd-ti-embedding' in embed_image.text: if hasattr(embed_image, 'text') and 'sd-ti-embedding' in embed_image.text:
data = embedding_from_b64(embed_image.text['sd-ti-embedding']) data = embedding_from_b64(embed_image.text['sd-ti-embedding'])
name = data.get('name', name) name = data.get('name', name)
else: else:
data = extract_image_data_embed(embed_image) data = extract_image_data_embed(embed_image)
name = data.get('name', name) name = data.get('name', name)
elif filename.upper().endswith('.WEBP'):
embed_image = Image.open(path)
data = extract_image_data_embed(embed_image)
name = data.get('name', name)
else: else:
data = torch.load(path, map_location="cpu") data = torch.load(path, map_location="cpu")