From 99bd6b0708303befd8caba1ee3b4488ef746871f Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Mon, 8 Jan 2024 22:32:49 -0600 Subject: [PATCH] keep safe images --- api/onnx_web/chain/edit_safety.py | 4 ++++ api/onnx_web/diffusers/run.py | 12 +++++++++--- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/api/onnx_web/chain/edit_safety.py b/api/onnx_web/chain/edit_safety.py index 35d535b5..0a1d281f 100644 --- a/api/onnx_web/chain/edit_safety.py +++ b/api/onnx_web/chain/edit_safety.py @@ -58,10 +58,14 @@ class EditSafetyStage(BaseStage): if check.is_csam: logger.warning("flagging csam result: %s, %s", i, prompt) is_csam = True + continue if check.is_nsfw and block_nsfw: logger.warning("blocking nsfw image: %s, %s", i, prompt) results.append(Image.new("RGB", image.size, color="black")) + continue + + results.append(image) if is_csam: logger.warning("blocking csam result") diff --git a/api/onnx_web/diffusers/run.py b/api/onnx_web/diffusers/run.py index ac6a36c4..0947c85f 100644 --- a/api/onnx_web/diffusers/run.py +++ b/api/onnx_web/diffusers/run.py @@ -132,10 +132,16 @@ def run_txt2img_pipeline( ) # add a thumbnail, if requested - cover = images.as_images()[0] - if params.thumbnail and ( - cover.width > server.thumbnail_size or cover.height > server.thumbnail_size + result_size = images.size() + if ( + params.thumbnail + and len(images) > 0 + and ( + result_size.width > server.thumbnail_size + or result_size.height > server.thumbnail_size + ) ): + cover = images.as_images()[0] thumbnail = cover.copy() thumbnail.thumbnail((server.thumbnail_size, server.thumbnail_size))