From 55e23a9374cbe09d70f182b7a73a7885411822db Mon Sep 17 00:00:00 2001 From: Spencer Hayes-Laverdiere Date: Fri, 15 Mar 2024 17:30:29 -0400 Subject: [PATCH] fix: add error output for unsupported images (#2537) * Raise Error on bad decode * Move task arg pop to try block * fix: prevent empty task from getting queued --------- Co-authored-by: Manuel Schmid --- modules/gradio_hijack.py | 7 +++++-- webui.py | 6 +++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/modules/gradio_hijack.py b/modules/gradio_hijack.py index 181429e..35df81c 100644 --- a/modules/gradio_hijack.py +++ b/modules/gradio_hijack.py @@ -17,7 +17,7 @@ from gradio_client.documentation import document, set_documentation_group from gradio_client.serializing import ImgSerializable from PIL import Image as _Image # using _ to minimize namespace pollution -from gradio import processing_utils, utils +from gradio import processing_utils, utils, Error from gradio.components.base import IOComponent, _Keywords, Block from gradio.deprecation import warn_style_method_deprecation from gradio.events import ( @@ -275,7 +275,10 @@ class Image( x, mask = x["image"], x["mask"] assert isinstance(x, str) - im = processing_utils.decode_base64_to_image(x) + try: + im = processing_utils.decode_base64_to_image(x) + except PIL.UnidentifiedImageError: + raise Error("Unsupported image type in input") with warnings.catch_warnings(): warnings.simplefilter("ignore") im = im.convert(self.image_mode) diff --git a/webui.py b/webui.py index 01c828d..98780bf 100644 --- a/webui.py +++ b/webui.py @@ -29,12 +29,16 @@ def get_task(*args): return worker.AsyncTask(args=args) -def generate_clicked(task): +def generate_clicked(task: worker.AsyncTask): import ldm_patched.modules.model_management as model_management with model_management.interrupt_processing_mutex: model_management.interrupt_processing = False # outputs=[progress_html, progress_window, progress_gallery, gallery] + + if len(task.args) == 0: + return + execution_start_time = time.perf_counter() finished = False