From 5de28d0361b66a522d965cf37f39bc8d52dbe56b Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Tue, 2 Jan 2024 20:58:02 -0600
Subject: [PATCH] generate a thumbnail when the first image is > 512

---
 api/onnx_web/diffusers/run.py | 7 +++++++
 api/onnx_web/params.py        | 5 +++++
 2 files changed, 12 insertions(+)

diff --git a/api/onnx_web/diffusers/run.py b/api/onnx_web/diffusers/run.py
index a3737d06..9e902cc2 100644
--- a/api/onnx_web/diffusers/run.py
+++ b/api/onnx_web/diffusers/run.py
@@ -120,6 +120,13 @@ def run_txt2img_pipeline(
 
     _pairs, loras, inversions, _rest = parse_prompt(params)
 
+    # add a thumbnail, if requested
+    cover = images[0]
+    if params.thumbnail and (cover.width > 512 or cover.height > 512):
+        thumbnail = cover.thumbnail((512, 512))
+        images.insert(thumbnail)
+        outputs.insert(f"{worker.name}-thumb.{server.image_format}")
+
     for image, output in zip(images, outputs):
         logger.trace("saving output image %s: %s", output, image.size)
         dest = save_image(
diff --git a/api/onnx_web/params.py b/api/onnx_web/params.py
index 8d255183..b8dc3049 100644
--- a/api/onnx_web/params.py
+++ b/api/onnx_web/params.py
@@ -213,6 +213,7 @@ class ImageParams:
     vae_tile: int
     vae_overlap: float
     denoise: int
+    thumbnail: int
 
     def __init__(
         self,
@@ -236,6 +237,7 @@ class ImageParams:
         vae_overlap: float = 0.25,
         vae_tile: int = 512,
         denoise: int = 3,
+        thumbnail: int = 1,
     ) -> None:
         self.model = model
         self.pipeline = pipeline
@@ -257,6 +259,7 @@ class ImageParams:
         self.vae_overlap = vae_overlap
         self.vae_tile = vae_tile
         self.denoise = denoise
+        self.thumbnail = thumbnail
 
     def do_cfg(self):
         return self.cfg > 1.0
@@ -328,6 +331,7 @@ class ImageParams:
             "vae_overlap": self.vae_overlap,
             "vae_tile": self.vae_tile,
             "denoise": self.denoise,
+            "thumbnail": self.thumbnail,
         }
 
     def with_args(self, **kwargs):
@@ -352,6 +356,7 @@ class ImageParams:
             kwargs.get("vae_overlap", self.vae_overlap),
             kwargs.get("vae_tile", self.vae_tile),
             kwargs.get("denoise", self.denoise),
+            kwargs.get("thumbnail", self.thumbnail),
         )