1
0
Fork 0

add DAT stage to compound upscale stage

This commit is contained in:
Sean Sube 2023-12-31 11:27:06 -06:00
parent 4a9ca4c4a8
commit 8c54de7a52
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 7 additions and 0 deletions

View File

@ -6,6 +6,7 @@ from . import ChainPipeline, PipelineStage
from .correct_codeformer import CorrectCodeformerStage from .correct_codeformer import CorrectCodeformerStage
from .correct_gfpgan import CorrectGFPGANStage from .correct_gfpgan import CorrectGFPGANStage
from .upscale_bsrgan import UpscaleBSRGANStage from .upscale_bsrgan import UpscaleBSRGANStage
from .upscale_dat import UpscaleDATStage
from .upscale_resrgan import UpscaleRealESRGANStage from .upscale_resrgan import UpscaleRealESRGANStage
from .upscale_stable_diffusion import UpscaleStableDiffusionStage from .upscale_stable_diffusion import UpscaleStableDiffusionStage
from .upscale_swinir import UpscaleSwinIRStage from .upscale_swinir import UpscaleSwinIRStage
@ -77,6 +78,12 @@ def stage_upscale_correction(
outscale=upscale.outscale, outscale=upscale.outscale,
) )
upscale_stage = (UpscaleBSRGANStage(), bsrgan_params, upscale_opts) 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: elif "esrgan" in upscale.upscale_model:
esrgan_params = StageParams( esrgan_params = StageParams(
tile_size=stage.tile_size, tile_size=stage.tile_size,