From 5ffb44c8fad51f9dfa2150c72106002d579a567e Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 14 Jan 2024 15:36:21 -0600 Subject: [PATCH] reset outscale for thumbnail metadata --- api/onnx_web/chain/edit_safety.py | 3 +- api/onnx_web/diffusers/run.py | 10 +++++ api/onnx_web/params.py | 70 +++++++++++++++++++++++-------- 3 files changed, 65 insertions(+), 18 deletions(-) diff --git a/api/onnx_web/chain/edit_safety.py b/api/onnx_web/chain/edit_safety.py index ee5eeb78..71650892 100644 --- a/api/onnx_web/chain/edit_safety.py +++ b/api/onnx_web/chain/edit_safety.py @@ -18,10 +18,12 @@ class EditSafetyStage(BaseStage): max_tile = SizeChart.max def load(self, server: ServerContext) -> Any: + # keep these within run to make this sort of like a plugin or peer dependency from horde_safety.deep_danbooru_model import get_deep_danbooru_model from horde_safety.interrogate import get_interrogator_no_blip from horde_safety.nsfw_checker_class import NSFWChecker + # check cache cache_key = ("horde-safety",) cache_checker = server.cache.get(ModelTypes.safety, cache_key) if cache_checker is not None: @@ -53,7 +55,6 @@ class EditSafetyStage(BaseStage): ) -> StageResult: logger.info("checking results using horde safety") - # keep these within run to make this sort of like a plugin or peer dependency try: # set up nsfw_checker = self.load(server) diff --git a/api/onnx_web/diffusers/run.py b/api/onnx_web/diffusers/run.py index 816aa01d..04e2f2f4 100644 --- a/api/onnx_web/diffusers/run.py +++ b/api/onnx_web/diffusers/run.py @@ -91,6 +91,16 @@ def add_thumbnail_output( size=Size(server.thumbnail_size, server.thumbnail_size) ) + if metadata.highres is not None: + metadata.highres = metadata.highres.with_args( + outscale=1, + ) + + if metadata.upscale is not None: + metadata.upscale = metadata.upscale.with_args( + outscale=1, + ) + images.insert_image(0, thumbnail, metadata) diff --git a/api/onnx_web/params.py b/api/onnx_web/params.py index 6662e067..c523750e 100644 --- a/api/onnx_web/params.py +++ b/api/onnx_web/params.py @@ -405,6 +405,9 @@ class StageParams: ) +UpscaleOrder = Literal["correction-first", "correction-last", "correction-both"] + + class UpscaleParams: def __init__( self, @@ -419,9 +422,7 @@ class UpscaleParams: scale: int = 4, pre_pad: int = 0, tile_pad: int = 10, - upscale_order: Literal[ - "correction-first", "correction-last", "correction-both" - ] = "correction-first", + upscale_order: UpscaleOrder = "correction-first", ) -> None: self.upscale_model = upscale_model self.correction_model = correction_model @@ -478,23 +479,40 @@ class UpscaleParams: "upscale_order": self.upscale_order, } - def with_args(self, **kwargs): + def with_args( + self, + upscale_model: Optional[str] = None, + correction_model: Optional[str] = None, + denoise: Optional[float] = None, + upscale: Optional[bool] = None, + faces: Optional[bool] = None, + face_outscale: Optional[int] = None, + face_strength: Optional[float] = None, + outscale: Optional[int] = None, + scale: Optional[int] = None, + pre_pad: Optional[int] = None, + tile_pad: Optional[int] = None, + upscale_order: Optional[UpscaleOrder] = None, + ): return UpscaleParams( - kwargs.get("upscale_model", self.upscale_model), - kwargs.get("correction_model", self.correction_model), - kwargs.get("denoise", self.denoise), - kwargs.get("upscale", self.upscale), - kwargs.get("faces", self.faces), - kwargs.get("face_outscale", self.face_outscale), - kwargs.get("face_strength", self.face_strength), - kwargs.get("outscale", self.outscale), - kwargs.get("scale", self.scale), - kwargs.get("pre_pad", self.pre_pad), - kwargs.get("tile_pad", self.tile_pad), - kwargs.get("upscale_order", self.upscale_order), + upscale_model=upscale_model or self.upscale_model, + correction_model=correction_model or self.correction_model, + denoise=denoise or self.denoise, + upscale=upscale or self.upscale, + faces=faces or self.faces, + face_outscale=face_outscale or self.face_outscale, + face_strength=face_strength or self.face_strength, + outscale=outscale or self.outscale, + scale=scale or self.scale, + pre_pad=pre_pad or self.pre_pad, + tile_pad=tile_pad or self.tile_pad, + upscale_order=upscale_order or self.upscale_order, ) +UpscaleMethod = Literal["bilinear", "lanczos", "upscale"] + + class HighresParams: def __init__( self, @@ -502,7 +520,7 @@ class HighresParams: scale: int, steps: int, strength: float, - method: Literal["bilinear", "lanczos", "upscale"] = "lanczos", + method: UpscaleMethod = "lanczos", iterations: int = 1, ): self.enabled = enabled @@ -531,3 +549,21 @@ class HighresParams: "steps": self.steps, "strength": self.strength, } + + def with_args( + self, + enabled: Optional[bool] = None, + scale: Optional[int] = None, + steps: Optional[int] = None, + strength: Optional[float] = None, + method: Optional[UpscaleMethod] = None, + iterations: Optional[int] = None, + ): + return HighresParams( + enabled=enabled or self.enabled, + scale=scale or self.scale, + steps=steps or self.steps, + strength=strength or self.strength, + method=method or self.method, + iterations=iterations or self.iterations, + )