From 95e2d6d71006b02b24b5ce72e6eb9078e4e0ffaa Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 12 Nov 2023 21:13:52 -0600 Subject: [PATCH] feat(api): add denoise stage, use before highres --- api/onnx_web/chain/__init__.py | 2 + api/onnx_web/chain/base.py | 5 ++- api/onnx_web/chain/blend_denoise.py | 39 +++++++++++++++++++ api/onnx_web/diffusers/pipelines/panorama.py | 2 +- .../diffusers/pipelines/panorama_xl.py | 2 +- api/onnx_web/diffusers/run.py | 9 ++++- 6 files changed, 54 insertions(+), 5 deletions(-) create mode 100644 api/onnx_web/chain/blend_denoise.py diff --git a/api/onnx_web/chain/__init__.py b/api/onnx_web/chain/__init__.py index df2ac80e..b34372ee 100644 --- a/api/onnx_web/chain/__init__.py +++ b/api/onnx_web/chain/__init__.py @@ -1,4 +1,5 @@ from .base import ChainPipeline, PipelineStage, StageParams +from .blend_denoise import BlendDenoiseStage from .blend_img2img import BlendImg2ImgStage from .blend_grid import BlendGridStage from .blend_linear import BlendLinearStage @@ -22,6 +23,7 @@ from .upscale_stable_diffusion import UpscaleStableDiffusionStage from .upscale_swinir import UpscaleSwinIRStage CHAIN_STAGES = { + "blend-denoise": BlendDenoiseStage, "blend-img2img": BlendImg2ImgStage, "blend-inpaint": UpscaleOutpaintStage, "blend-grid": BlendGridStage, diff --git a/api/onnx_web/chain/base.py b/api/onnx_web/chain/base.py index 79323081..cabb8da1 100644 --- a/api/onnx_web/chain/base.py +++ b/api/onnx_web/chain/base.py @@ -232,7 +232,10 @@ class ChainPipeline: stage_sources = stage_outputs else: - logger.debug("image does not contain sources and is within tile size of %s, running stage", tile) + logger.debug( + "image does not contain sources and is within tile size of %s, running stage", + tile, + ) for i in range(worker.retries): try: stage_outputs = stage_pipe.run( diff --git a/api/onnx_web/chain/blend_denoise.py b/api/onnx_web/chain/blend_denoise.py new file mode 100644 index 00000000..9c23c609 --- /dev/null +++ b/api/onnx_web/chain/blend_denoise.py @@ -0,0 +1,39 @@ +from logging import getLogger +from typing import List, Optional + +import cv2 +import numpy as np +from PIL import Image + +from ..params import ImageParams, SizeChart, StageParams +from ..server import ServerContext +from ..worker import ProgressCallback, WorkerContext +from .stage import BaseStage + +logger = getLogger(__name__) + + +class BlendDenoiseStage(BaseStage): + max_tile = SizeChart.max + + def run( + self, + _worker: WorkerContext, + _server: ServerContext, + _stage: StageParams, + _params: ImageParams, + sources: List[Image.Image], + *, + stage_source: Optional[Image.Image] = None, + callback: Optional[ProgressCallback] = None, + **kwargs, + ) -> List[Image.Image]: + logger.info("denoising source images") + + results = [] + for source in sources: + data = cv2.cvtColor(np.array(source), cv2.COLOR_RGB2BGR) + data = cv2.fastNlMeansDenoisingColored(data) + results.append(Image.fromarray(cv2.cvtColor(data, cv2.COLOR_BGR2RGB))) + + return results diff --git a/api/onnx_web/diffusers/pipelines/panorama.py b/api/onnx_web/diffusers/pipelines/panorama.py index dd737a1e..01d6d0ec 100644 --- a/api/onnx_web/diffusers/pipelines/panorama.py +++ b/api/onnx_web/diffusers/pipelines/panorama.py @@ -702,12 +702,12 @@ class OnnxStableDiffusionPanoramaPipeline(DiffusionPipeline): # take the MultiDiffusion step. Eq. 5 in MultiDiffusion paper: https://arxiv.org/abs/2302.08113 latents = np.where(count > 0, value / count, value) latents = repair_nan(latents) - latents = np.clip(latents, -4, +4) # call the callback, if provided if callback is not None and i % callback_steps == 0: callback(i, t, latents) + latents = np.clip(latents, -4, +4) latents = 1 / 0.18215 * latents # image = self.vae_decoder(latent_sample=latents)[0] # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 diff --git a/api/onnx_web/diffusers/pipelines/panorama_xl.py b/api/onnx_web/diffusers/pipelines/panorama_xl.py index e0e26621..446bc5d6 100644 --- a/api/onnx_web/diffusers/pipelines/panorama_xl.py +++ b/api/onnx_web/diffusers/pipelines/panorama_xl.py @@ -551,7 +551,6 @@ class StableDiffusionXLPanoramaPipelineMixin(StableDiffusionXLImg2ImgPipelineMix # take the MultiDiffusion step. Eq. 5 in MultiDiffusion paper: https://arxiv.org/abs/2302.08113 latents = np.where(count > 0, value / count, value) latents = repair_nan(latents) - latents = np.clip(latents, -4, +4) # call the callback, if provided if i == len(timesteps) - 1 or ( @@ -563,6 +562,7 @@ class StableDiffusionXLPanoramaPipelineMixin(StableDiffusionXLImg2ImgPipelineMix if output_type == "latent": image = latents else: + latents = np.clip(latents, -4, +4) latents = latents / self.vae_decoder.config.get("scaling_factor", 0.18215) # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 image = np.concatenate( diff --git a/api/onnx_web/diffusers/run.py b/api/onnx_web/diffusers/run.py index a18cb75d..a3cd66fe 100644 --- a/api/onnx_web/diffusers/run.py +++ b/api/onnx_web/diffusers/run.py @@ -4,15 +4,15 @@ from typing import Any, List, Optional from PIL import Image, ImageOps -from onnx_web.chain.highres import stage_highres - from ..chain import ( + BlendDenoiseStage, BlendImg2ImgStage, BlendMaskStage, ChainPipeline, SourceTxt2ImgStage, UpscaleOutpaintStage, ) +from ..chain.highres import stage_highres from ..chain.upscale import split_upscale, stage_upscale_correction from ..image import expand_image from ..output import save_image @@ -68,6 +68,11 @@ def run_txt2img_pipeline( highres_size = params.unet_tile stage = StageParams(tile_size=highres_size) + chain.stage( + BlendDenoiseStage(), + stage, + ) + first_upscale, after_upscale = split_upscale(upscale) if first_upscale: stage_upscale_correction(