From 8c54de7a5264de82b21ae0e0f6da0606ee545d39 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 31 Dec 2023 11:27:06 -0600 Subject: [PATCH] add DAT stage to compound upscale stage --- api/onnx_web/chain/upscale.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/api/onnx_web/chain/upscale.py b/api/onnx_web/chain/upscale.py index db622c36..23d1e065 100644 --- a/api/onnx_web/chain/upscale.py +++ b/api/onnx_web/chain/upscale.py @@ -6,6 +6,7 @@ from . import ChainPipeline, PipelineStage from .correct_codeformer import CorrectCodeformerStage from .correct_gfpgan import CorrectGFPGANStage from .upscale_bsrgan import UpscaleBSRGANStage +from .upscale_dat import UpscaleDATStage from .upscale_resrgan import UpscaleRealESRGANStage from .upscale_stable_diffusion import UpscaleStableDiffusionStage from .upscale_swinir import UpscaleSwinIRStage @@ -77,6 +78,12 @@ def stage_upscale_correction( outscale=upscale.outscale, ) upscale_stage = (UpscaleBSRGANStage(), bsrgan_params, upscale_opts) + elif "dat" in upscale.upscale_model: + dat_params = StageParams( + tile_size=stage.tile_size, + outscale=upscale.outscale, + ) + upscale_stage = (UpscaleDATStage(), dat_params, upscale_opts) elif "esrgan" in upscale.upscale_model: esrgan_params = StageParams( tile_size=stage.tile_size,