Fixes: uploaded JPG images not loading

This commit is contained in:
psychedelicious 2022-11-02 16:14:50 +11:00 committed by Lincoln Stein
parent cefe12f1df
commit cfe567c62a

View File

@ -187,7 +187,10 @@ class InvokeAIWebServer:
base_path = ( base_path = (
self.result_path if category == "result" else self.init_image_path self.result_path if category == "result" else self.init_image_path
) )
paths = glob.glob(os.path.join(base_path, "*.png"))
paths = []
for ext in ("*.png", "*.jpg", "*.jpeg"):
paths.extend(glob.glob(os.path.join(base_path, ext)))
image_paths = sorted( image_paths = sorted(
paths, key=lambda x: os.path.getmtime(x), reverse=True paths, key=lambda x: os.path.getmtime(x), reverse=True
@ -203,13 +206,19 @@ class InvokeAIWebServer:
image_array = [] image_array = []
for path in image_paths: for path in image_paths:
if os.path.splitext(path)[1] == ".png":
metadata = retrieve_metadata(path) metadata = retrieve_metadata(path)
sd_metadata = metadata["sd-metadata"]
else:
sd_metadata = {}
(width, height) = Image.open(path).size (width, height) = Image.open(path).size
image_array.append( image_array.append(
{ {
"url": self.get_url_from_image_path(path), "url": self.get_url_from_image_path(path),
"mtime": os.path.getmtime(path), "mtime": os.path.getmtime(path),
"metadata": metadata["sd-metadata"], "metadata": sd_metadata,
"width": width, "width": width,
"height": height, "height": height,
"category": category, "category": category,
@ -236,7 +245,9 @@ class InvokeAIWebServer:
self.result_path if category == "result" else self.init_image_path self.result_path if category == "result" else self.init_image_path
) )
paths = glob.glob(os.path.join(base_path, "*.png")) paths = []
for ext in ("*.png", "*.jpg", "*.jpeg"):
paths.extend(glob.glob(os.path.join(base_path, ext)))
image_paths = sorted( image_paths = sorted(
paths, key=lambda x: os.path.getmtime(x), reverse=True paths, key=lambda x: os.path.getmtime(x), reverse=True
@ -254,9 +265,12 @@ class InvokeAIWebServer:
image_paths = image_paths[slice(0, page_size)] image_paths = image_paths[slice(0, page_size)]
image_array = [] image_array = []
for path in image_paths: for path in image_paths:
if os.path.splitext(path)[1] == ".png":
metadata = retrieve_metadata(path) metadata = retrieve_metadata(path)
sd_metadata = metadata["sd-metadata"]
else:
sd_metadata = {}
(width, height) = Image.open(path).size (width, height) = Image.open(path).size
@ -264,7 +278,7 @@ class InvokeAIWebServer:
{ {
"url": self.get_url_from_image_path(path), "url": self.get_url_from_image_path(path),
"mtime": os.path.getmtime(path), "mtime": os.path.getmtime(path),
"metadata": metadata["sd-metadata"], "metadata": sd_metadata,
"width": width, "width": width,
"height": height, "height": height,
"category": category, "category": category,
@ -601,8 +615,9 @@ class InvokeAIWebServer:
progress.set_current_status_has_steps(True) progress.set_current_status_has_steps(True)
if ( if (
generation_parameters['progress_images'] and step % 5 == 0 \ generation_parameters["progress_images"]
and step < generation_parameters['steps'] - 1 and step % 5 == 0
and step < generation_parameters["steps"] - 1
): ):
image = self.generate.sample_to_image(sample) image = self.generate.sample_to_image(sample)
metadata = self.parameters_to_generated_image_metadata( metadata = self.parameters_to_generated_image_metadata(
@ -633,14 +648,16 @@ class InvokeAIWebServer:
}, },
) )
if generation_parameters['progress_latents']: if generation_parameters["progress_latents"]:
image = self.generate.sample_to_lowres_estimated_image(sample) image = self.generate.sample_to_lowres_estimated_image(sample)
(width, height) = image.size (width, height) = image.size
width *= 8 width *= 8
height *= 8 height *= 8
buffered = io.BytesIO() buffered = io.BytesIO()
image.save(buffered, format="PNG") image.save(buffered, format="PNG")
img_base64 = "data:image/png;base64," + base64.b64encode(buffered.getvalue()).decode('UTF-8') img_base64 = "data:image/png;base64," + base64.b64encode(
buffered.getvalue()
).decode("UTF-8")
self.socketio.emit( self.socketio.emit(
"intermediateResult", "intermediateResult",
{ {
@ -650,7 +667,7 @@ class InvokeAIWebServer:
"metadata": {}, "metadata": {},
"width": width, "width": width,
"height": height, "height": height,
} },
) )
self.socketio.emit("progressUpdate", progress.to_formatted_dict()) self.socketio.emit("progressUpdate", progress.to_formatted_dict())
@ -668,7 +685,6 @@ class InvokeAIWebServer:
step_index = 1 step_index = 1
nonlocal prior_variations nonlocal prior_variations
# paste the inpainting image back onto the original # paste the inpainting image back onto the original
if "init_mask" in generation_parameters: if "init_mask" in generation_parameters:
image = paste_image_into_bounding_box( image = paste_image_into_bounding_box(