Gracefully handles corrupted images; fixes #1486

- App does not crash if corrupted image loaded
- Error is displayed in the UI console and CLI output if an image cannot be loaded
This commit is contained in:
psychedelicious 2022-11-20 11:16:41 +11:00 committed by blessedcoolant
parent b908f2b4bc
commit 9d34213b4c

View File

@ -334,32 +334,34 @@ class InvokeAIWebServer:
image_array = [] image_array = []
for path in image_paths: for path in image_paths:
if os.path.splitext(path)[1] == ".png": try:
metadata = retrieve_metadata(path) if os.path.splitext(path)[1] == ".png":
# sd_metadata = metadata["sd-metadata"] metadata = retrieve_metadata(path)
else: else:
# sd_metadata = {} metadata = {}
metadata = {}
pil_image = Image.open(path) pil_image = Image.open(path)
(width, height) = pil_image.size (width, height) = pil_image.size
thumbnail_path = save_thumbnail( thumbnail_path = save_thumbnail(
pil_image, os.path.basename(path), self.thumbnail_image_path pil_image, os.path.basename(path), self.thumbnail_image_path
) )
image_array.append( image_array.append(
{ {
"url": self.get_url_from_image_path(path), "url": self.get_url_from_image_path(path),
"thumbnail": self.get_url_from_image_path(thumbnail_path), "thumbnail": self.get_url_from_image_path(
"mtime": os.path.getmtime(path), thumbnail_path
"metadata": metadata, ),
# "metadata": sd_metadata, "mtime": os.path.getmtime(path),
"width": width, "metadata": metadata,
"height": height, "width": width,
"category": category, "height": height,
} "category": category,
) }
)
except:
socketio.emit("error", {"message": f"Unable to load {path}"})
socketio.emit( socketio.emit(
"galleryImages", "galleryImages",
@ -402,32 +404,35 @@ class InvokeAIWebServer:
image_array = [] image_array = []
for path in image_paths: for path in image_paths:
if os.path.splitext(path)[1] == ".png": try:
metadata = retrieve_metadata(path) if os.path.splitext(path)[1] == ".png":
# sd_metadata = metadata["sd-metadata"] metadata = retrieve_metadata(path)
else: else:
# sd_metadata = {} metadata = {}
metadata = {}
pil_image = Image.open(path) pil_image = Image.open(path)
(width, height) = pil_image.size (width, height) = pil_image.size
thumbnail_path = save_thumbnail( thumbnail_path = save_thumbnail(
pil_image, os.path.basename(path), self.thumbnail_image_path pil_image, os.path.basename(path), self.thumbnail_image_path
) )
image_array.append( image_array.append(
{ {
"url": self.get_url_from_image_path(path), "url": self.get_url_from_image_path(path),
"thumbnail": self.get_url_from_image_path(thumbnail_path), "thumbnail": self.get_url_from_image_path(
"mtime": os.path.getmtime(path), thumbnail_path
# "metadata": sd_metadata, ),
"metadata": metadata, "mtime": os.path.getmtime(path),
"width": width, "metadata": metadata,
"height": height, "width": width,
"category": category, "height": height,
} "category": category,
) }
)
except:
print(f'>> Unable to load {path}')
socketio.emit("error", {"message": f"Unable to load {path}"})
socketio.emit( socketio.emit(
"galleryImages", "galleryImages",