2023-01-16 00:04:10 +00:00
|
|
|
from basicsr.archs.rrdbnet_arch import RRDBNet
|
2023-01-26 03:29:18 +00:00
|
|
|
from diffusers import (
|
|
|
|
AutoencoderKL,
|
|
|
|
DDPMScheduler,
|
2023-01-27 04:44:20 +00:00
|
|
|
StableDiffusionUpscalePipeline,
|
2023-01-26 03:29:18 +00:00
|
|
|
)
|
2023-01-16 00:04:10 +00:00
|
|
|
from gfpgan import GFPGANer
|
|
|
|
from os import path
|
|
|
|
from PIL import Image
|
|
|
|
from realesrgan import RealESRGANer
|
2023-01-27 23:08:36 +00:00
|
|
|
from typing import Literal, Optional
|
2023-01-16 00:04:10 +00:00
|
|
|
|
2023-01-16 00:13:28 +00:00
|
|
|
import numpy as np
|
2023-01-27 05:15:24 +00:00
|
|
|
import torch
|
2023-01-16 00:13:28 +00:00
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
from .chain import (
|
|
|
|
ChainPipeline,
|
|
|
|
StageParams,
|
2023-01-27 04:44:20 +00:00
|
|
|
)
|
2023-01-26 03:04:00 +00:00
|
|
|
from .onnx import (
|
|
|
|
ONNXNet,
|
|
|
|
OnnxStableDiffusionUpscalePipeline,
|
|
|
|
)
|
2023-01-16 19:02:15 +00:00
|
|
|
from .utils import (
|
2023-01-27 23:08:36 +00:00
|
|
|
ImageParams,
|
2023-01-16 21:11:40 +00:00
|
|
|
ServerContext,
|
|
|
|
Size,
|
2023-01-16 19:02:15 +00:00
|
|
|
)
|
|
|
|
|
2023-01-16 16:55:40 +00:00
|
|
|
|
2023-01-16 19:02:15 +00:00
|
|
|
class UpscaleParams():
|
2023-01-16 20:52:56 +00:00
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
upscale_model: str,
|
2023-01-22 22:35:53 +00:00
|
|
|
provider: str,
|
2023-01-27 23:08:36 +00:00
|
|
|
correction_model: Optional[str] = None,
|
2023-01-16 20:52:56 +00:00
|
|
|
denoise: float = 0.5,
|
|
|
|
faces=True,
|
2023-01-17 04:54:01 +00:00
|
|
|
face_strength: float = 0.5,
|
2023-01-22 22:35:53 +00:00
|
|
|
format: Literal['onnx', 'pth'] = 'onnx',
|
|
|
|
half=False,
|
|
|
|
outscale: int = 1,
|
|
|
|
scale: int = 4,
|
2023-01-27 23:08:36 +00:00
|
|
|
pre_pad: int = 0,
|
|
|
|
tile_pad: int = 10,
|
2023-01-16 20:52:56 +00:00
|
|
|
) -> None:
|
|
|
|
self.upscale_model = upscale_model
|
2023-01-22 22:35:53 +00:00
|
|
|
self.provider = provider
|
2023-01-17 02:10:52 +00:00
|
|
|
self.correction_model = correction_model
|
2023-01-16 20:52:56 +00:00
|
|
|
self.denoise = denoise
|
2023-01-16 19:02:15 +00:00
|
|
|
self.faces = faces
|
2023-01-17 04:54:01 +00:00
|
|
|
self.face_strength = face_strength
|
2023-01-22 22:35:53 +00:00
|
|
|
self.format = format
|
2023-01-16 20:52:56 +00:00
|
|
|
self.half = half
|
2023-01-22 22:35:53 +00:00
|
|
|
self.outscale = outscale
|
2023-01-27 23:08:36 +00:00
|
|
|
self.pre_pad = pre_pad
|
2023-01-22 22:35:53 +00:00
|
|
|
self.scale = scale
|
2023-01-27 23:08:36 +00:00
|
|
|
self.tile_pad = tile_pad
|
2023-01-16 19:02:15 +00:00
|
|
|
|
2023-01-16 21:11:40 +00:00
|
|
|
def resize(self, size: Size) -> Size:
|
2023-01-20 01:42:27 +00:00
|
|
|
return Size(size.width * self.outscale, size.height * self.outscale)
|
2023-01-16 21:11:40 +00:00
|
|
|
|
2023-01-16 19:02:15 +00:00
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
def load_resrgan(ctx: ServerContext, params: UpscaleParams, tile=0):
|
|
|
|
'''
|
|
|
|
TODO: cache this instance
|
|
|
|
'''
|
2023-01-22 22:35:53 +00:00
|
|
|
model_file = '%s.%s' % (params.upscale_model, params.format)
|
2023-01-16 20:58:27 +00:00
|
|
|
model_path = path.join(ctx.model_path, model_file)
|
2023-01-16 00:04:10 +00:00
|
|
|
if not path.isfile(model_path):
|
2023-01-16 20:52:56 +00:00
|
|
|
raise Exception('Real ESRGAN model not found at %s' % model_path)
|
2023-01-16 00:04:10 +00:00
|
|
|
|
2023-01-16 19:02:15 +00:00
|
|
|
# use ONNX acceleration, if available
|
2023-01-22 22:35:53 +00:00
|
|
|
if params.format == 'onnx':
|
|
|
|
model = ONNXNet(ctx, model_file, provider=params.provider)
|
|
|
|
elif params.format == 'pth':
|
2023-01-16 19:02:15 +00:00
|
|
|
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
|
|
|
|
num_block=23, num_grow_ch=32, scale=params.scale)
|
2023-01-22 22:35:53 +00:00
|
|
|
raise Exception('unknown platform %s' % params.format)
|
2023-01-16 00:04:10 +00:00
|
|
|
|
|
|
|
dni_weight = None
|
2023-01-16 20:52:56 +00:00
|
|
|
if params.upscale_model == 'realesr-general-x4v3' and params.denoise != 1:
|
2023-01-16 00:04:10 +00:00
|
|
|
wdn_model_path = model_path.replace(
|
|
|
|
'realesr-general-x4v3', 'realesr-general-wdn-x4v3')
|
|
|
|
model_path = [model_path, wdn_model_path]
|
2023-01-16 19:02:15 +00:00
|
|
|
dni_weight = [params.denoise, 1 - params.denoise]
|
2023-01-16 00:04:10 +00:00
|
|
|
|
2023-01-16 21:11:40 +00:00
|
|
|
# TODO: shouldn't need the PTH file
|
2023-01-16 00:04:10 +00:00
|
|
|
upsampler = RealESRGANer(
|
2023-01-16 19:02:15 +00:00
|
|
|
scale=params.scale,
|
2023-01-16 21:11:40 +00:00
|
|
|
model_path=path.join(ctx.model_path, '%s.pth' % params.upscale_model),
|
2023-01-16 00:04:10 +00:00
|
|
|
dni_weight=dni_weight,
|
|
|
|
model=model,
|
|
|
|
tile=tile,
|
2023-01-27 23:08:36 +00:00
|
|
|
tile_pad=params.tile_pad,
|
|
|
|
pre_pad=params.pre_pad,
|
2023-01-16 20:52:56 +00:00
|
|
|
half=params.half)
|
2023-01-16 00:04:10 +00:00
|
|
|
|
2023-01-16 00:41:58 +00:00
|
|
|
return upsampler
|
|
|
|
|
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
def load_stable_diffusion(ctx: ServerContext, upscale: UpscaleParams):
|
|
|
|
'''
|
|
|
|
TODO: cache this instance
|
|
|
|
'''
|
|
|
|
if upscale.format == 'onnx':
|
|
|
|
model_path = path.join(ctx.model_path, upscale.upscale_model)
|
|
|
|
# ValueError: Pipeline <class 'onnx_web.onnx.pipeline_onnx_stable_diffusion_upscale.OnnxStableDiffusionUpscalePipeline'>
|
|
|
|
# expected {'vae', 'unet', 'text_encoder', 'tokenizer', 'scheduler', 'low_res_scheduler'},
|
|
|
|
# but only {'scheduler', 'tokenizer', 'text_encoder', 'unet'} were passed.
|
|
|
|
pipeline = OnnxStableDiffusionUpscalePipeline.from_pretrained(
|
|
|
|
model_path,
|
|
|
|
vae=AutoencoderKL.from_pretrained(
|
|
|
|
model_path, subfolder='vae_encoder'),
|
|
|
|
low_res_scheduler=DDPMScheduler.from_pretrained(
|
|
|
|
model_path, subfolder='scheduler'),
|
|
|
|
)
|
|
|
|
else:
|
|
|
|
pipeline = StableDiffusionUpscalePipeline.from_pretrained(
|
|
|
|
'stabilityai/stable-diffusion-x4-upscaler')
|
|
|
|
|
|
|
|
return pipeline
|
|
|
|
|
|
|
|
|
|
|
|
def upscale_resrgan(
|
|
|
|
ctx: ServerContext,
|
|
|
|
stage: StageParams,
|
|
|
|
params: ImageParams,
|
|
|
|
source_image: Image.Image,
|
|
|
|
*,
|
|
|
|
upscale: UpscaleParams,
|
|
|
|
) -> Image:
|
|
|
|
print('upscaling image with Real ESRGAN', upscale.scale)
|
2023-01-16 19:12:08 +00:00
|
|
|
|
2023-01-17 05:01:15 +00:00
|
|
|
output = np.array(source_image)
|
2023-01-27 23:08:36 +00:00
|
|
|
upsampler = load_resrgan(ctx, upscale, tile=stage.tile_size)
|
2023-01-16 00:41:58 +00:00
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
output, _ = upsampler.enhance(output, outscale=upscale.outscale)
|
2023-01-16 00:25:12 +00:00
|
|
|
|
2023-01-18 04:30:25 +00:00
|
|
|
output = Image.fromarray(output, 'RGB')
|
|
|
|
print('final output image size', output.size)
|
|
|
|
return output
|
2023-01-16 00:04:10 +00:00
|
|
|
|
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
def upscale_gfpgan(
|
|
|
|
ctx: ServerContext,
|
|
|
|
stage: StageParams,
|
|
|
|
params: ImageParams,
|
|
|
|
image: Image.Image,
|
|
|
|
*,
|
|
|
|
upscale: UpscaleParams,
|
|
|
|
upsampler: Optional[RealESRGANer] = None,
|
|
|
|
) -> Image:
|
|
|
|
if upscale.correction_model is None:
|
2023-01-16 20:52:56 +00:00
|
|
|
print('no face model given, skipping')
|
|
|
|
return image
|
2023-01-16 19:12:08 +00:00
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
print('correcting faces with GFPGAN model: %s' % upscale.correction_model)
|
|
|
|
|
2023-01-16 19:02:15 +00:00
|
|
|
if upsampler is None:
|
2023-01-27 23:08:36 +00:00
|
|
|
upsampler = load_resrgan(ctx, upscale)
|
2023-01-16 19:02:15 +00:00
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
face_path = path.join(ctx.model_path, '%s.pth' %
|
|
|
|
(upscale.correction_model))
|
2023-01-16 22:45:17 +00:00
|
|
|
|
2023-01-16 23:48:50 +00:00
|
|
|
# TODO: doesn't have a model param, not sure how to pass ONNX model
|
2023-01-16 00:04:10 +00:00
|
|
|
face_enhancer = GFPGANer(
|
2023-01-16 22:45:17 +00:00
|
|
|
model_path=face_path,
|
2023-01-27 23:08:36 +00:00
|
|
|
upscale=upscale.outscale,
|
2023-01-16 00:04:10 +00:00
|
|
|
arch='clean',
|
|
|
|
channel_multiplier=2,
|
|
|
|
bg_upsampler=upsampler)
|
|
|
|
|
2023-01-16 00:41:58 +00:00
|
|
|
_, _, output = face_enhancer.enhance(
|
2023-01-27 23:08:36 +00:00
|
|
|
image, has_aligned=False, only_center_face=False, paste_back=True, weight=upscale.face_strength)
|
2023-01-16 00:04:10 +00:00
|
|
|
|
|
|
|
return output
|
2023-01-26 02:31:39 +00:00
|
|
|
|
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
def upscale_stable_diffusion(
|
|
|
|
ctx: ServerContext,
|
|
|
|
stage: StageParams,
|
|
|
|
params: ImageParams,
|
|
|
|
source: Image.Image,
|
|
|
|
*,
|
|
|
|
upscale: UpscaleParams,
|
|
|
|
) -> Image:
|
2023-01-26 03:04:00 +00:00
|
|
|
print('upscaling with Stable Diffusion')
|
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
pipeline = load_stable_diffusion(ctx, upscale)
|
|
|
|
generator = torch.manual_seed(params.seed)
|
|
|
|
seed = generator.initial_seed()
|
2023-01-26 03:04:00 +00:00
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
def upscale_stage(_ctx: ServerContext, stage: StageParams, params: ImageParams, image: Image.Image) -> Image:
|
|
|
|
return pipeline(
|
|
|
|
params.prompt,
|
2023-01-27 23:26:48 +00:00
|
|
|
image,
|
2023-01-27 23:08:36 +00:00
|
|
|
generator=torch.manual_seed(seed),
|
|
|
|
num_inference_steps=params.steps,
|
|
|
|
).images[0]
|
|
|
|
|
|
|
|
chain = ChainPipeline(stages=[
|
|
|
|
(upscale_stage, stage)
|
|
|
|
])
|
|
|
|
return chain(ctx, params, source)
|
|
|
|
|
|
|
|
|
|
|
|
def run_upscale_correction(
|
|
|
|
ctx: ServerContext,
|
|
|
|
stage: StageParams,
|
|
|
|
params: ImageParams,
|
|
|
|
image: Image.Image,
|
|
|
|
*,
|
|
|
|
upscale: UpscaleParams,
|
|
|
|
) -> Image.Image:
|
2023-01-26 02:31:39 +00:00
|
|
|
print('running upscale pipeline')
|
|
|
|
|
2023-01-27 23:38:21 +00:00
|
|
|
chain = ChainPipeline()
|
|
|
|
|
2023-01-27 23:08:36 +00:00
|
|
|
if upscale.scale > 1:
|
|
|
|
if 'esrgan' in upscale.upscale_model:
|
|
|
|
stage = StageParams(tile_size=stage.tile_size,
|
|
|
|
outscale=upscale.outscale)
|
2023-01-27 23:38:21 +00:00
|
|
|
chain.append((upscale_resrgan, stage, {'upscale': upscale}))
|
2023-01-27 23:08:36 +00:00
|
|
|
elif 'stable-diffusion' in upscale.upscale_model:
|
2023-01-27 23:26:48 +00:00
|
|
|
mini_tile = min(128, stage.tile_size)
|
2023-01-27 23:08:36 +00:00
|
|
|
stage = StageParams(tile_size=mini_tile, outscale=upscale.outscale)
|
2023-01-27 23:38:21 +00:00
|
|
|
chain.append((upscale_stable_diffusion, stage, {'upscale': upscale}))
|
2023-01-27 23:08:36 +00:00
|
|
|
|
|
|
|
if upscale.faces:
|
|
|
|
stage = StageParams(tile_size=stage.tile_size,
|
|
|
|
outscale=upscale.outscale)
|
2023-01-27 23:38:21 +00:00
|
|
|
chain.append((upscale_gfpgan, stage, {'upscale': upscale}))
|
2023-01-26 02:31:39 +00:00
|
|
|
|
2023-01-27 23:38:21 +00:00
|
|
|
return chain(ctx, params, image)
|