apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

keep safe images

This commit is contained in:
Sean Sube 2024-01-08 22:32:49 -06:00
parent 2eb90ba559
commit 99bd6b0708
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
2 changed files with 13 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
api/onnx_web/chain/edit_safety.py
View File

@ -58,10 +58,14 @@ class EditSafetyStage(BaseStage):
if check.is_csam: if check.is_csam:
logger.warning("flagging csam result: %s, %s", i, prompt) logger.warning("flagging csam result: %s, %s", i, prompt)
is_csam = True is_csam = True
continue
if check.is_nsfw and block_nsfw: if check.is_nsfw and block_nsfw:
logger.warning("blocking nsfw image: %s, %s", i, prompt) logger.warning("blocking nsfw image: %s, %s", i, prompt)
results.append(Image.new("RGB", image.size, color="black")) results.append(Image.new("RGB", image.size, color="black"))
continue
results.append(image)
if is_csam: if is_csam:
logger.warning("blocking csam result") logger.warning("blocking csam result")

12
api/onnx_web/diffusers/run.py
View File

@ -132,10 +132,16 @@ def run_txt2img_pipeline(
) )
# add a thumbnail, if requested # add a thumbnail, if requested
cover = images.as_images()[0] result_size = images.size()
if params.thumbnail and ( if (
cover.width > server.thumbnail_size or cover.height > server.thumbnail_size 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 = cover.copy()
thumbnail.thumbnail((server.thumbnail_size, server.thumbnail_size)) thumbnail.thumbnail((server.thumbnail_size, server.thumbnail_size))