2023-01-28 23:09:19 +00:00
|
|
|
from logging import getLogger
|
2023-01-28 05:28:14 +00:00
|
|
|
from os import path
|
2023-11-19 00:13:13 +00:00
|
|
|
from typing import Optional
|
2023-02-05 13:53:26 +00:00
|
|
|
|
2023-01-28 05:28:14 +00:00
|
|
|
from PIL import Image
|
|
|
|
|
2023-04-11 00:02:12 +00:00
|
|
|
from ..onnx import OnnxRRDBNet
|
2023-02-05 13:53:26 +00:00
|
|
|
from ..params import DeviceParams, ImageParams, StageParams, UpscaleParams
|
2023-07-03 16:33:56 +00:00
|
|
|
from ..server import ModelTypes, ServerContext
|
2023-02-19 02:28:21 +00:00
|
|
|
from ..utils import run_gc
|
2023-02-26 20:15:30 +00:00
|
|
|
from ..worker import WorkerContext
|
2023-11-18 23:18:23 +00:00
|
|
|
from .base import BaseStage
|
2023-11-19 00:08:38 +00:00
|
|
|
from .result import StageResult
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-01-28 23:09:19 +00:00
|
|
|
logger = getLogger(__name__)
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-02-18 14:41:15 +00:00
|
|
|
TAG_X4_V3 = "real-esrgan-x4-v3"
|
2023-02-07 03:36:20 +00:00
|
|
|
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-07-02 23:21:21 +00:00
|
|
|
class UpscaleRealESRGANStage(BaseStage):
|
2023-07-01 12:10:53 +00:00
|
|
|
def load(
|
|
|
|
self, server: ServerContext, params: UpscaleParams, device: DeviceParams, tile=0
|
|
|
|
):
|
|
|
|
# must be within load function for patches to take effect
|
|
|
|
# TODO: rewrite and remove
|
|
|
|
from realesrgan import RealESRGANer
|
|
|
|
|
2023-12-27 14:47:06 +00:00
|
|
|
class RealESRGANWrapper(RealESRGANer):
|
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
scale,
|
|
|
|
model_path,
|
|
|
|
dni_weight=None,
|
|
|
|
model=None,
|
|
|
|
tile=0,
|
|
|
|
tile_pad=10,
|
|
|
|
pre_pad=10,
|
|
|
|
half=False,
|
|
|
|
device=None,
|
|
|
|
gpu_id=None,
|
|
|
|
):
|
|
|
|
self.scale = scale
|
|
|
|
self.tile_size = tile
|
|
|
|
self.tile_pad = tile_pad
|
|
|
|
self.pre_pad = pre_pad
|
|
|
|
self.mod_scale = None
|
|
|
|
self.half = half
|
|
|
|
self.model = model
|
2023-12-27 15:01:38 +00:00
|
|
|
self.device = device
|
2023-12-27 14:47:06 +00:00
|
|
|
|
2024-01-06 02:13:16 +00:00
|
|
|
model_file = f"{params.upscale_model}.onnx"
|
2023-07-01 12:10:53 +00:00
|
|
|
model_path = path.join(server.model_path, model_file)
|
|
|
|
|
2024-01-06 02:13:16 +00:00
|
|
|
cache_key = (model_path, params.scale)
|
2023-07-03 16:33:56 +00:00
|
|
|
cache_pipe = server.cache.get(ModelTypes.upscaling, cache_key)
|
2023-07-01 12:10:53 +00:00
|
|
|
if cache_pipe is not None:
|
|
|
|
logger.info("reusing existing Real ESRGAN pipeline")
|
|
|
|
return cache_pipe
|
|
|
|
|
|
|
|
if not path.isfile(model_path):
|
|
|
|
raise FileNotFoundError("Real ESRGAN model not found at %s" % model_path)
|
|
|
|
|
|
|
|
# TODO: swap for regular RRDBNet after rewriting wrapper
|
|
|
|
model = OnnxRRDBNet(
|
|
|
|
server,
|
|
|
|
model_file,
|
2024-01-13 00:58:26 +00:00
|
|
|
provider=device.ort_provider("esrgan"),
|
2023-07-01 12:10:53 +00:00
|
|
|
sess_options=device.sess_options(),
|
|
|
|
)
|
|
|
|
|
|
|
|
dni_weight = None
|
|
|
|
if params.upscale_model == TAG_X4_V3 and params.denoise != 1:
|
|
|
|
wdn_model_path = model_path.replace(TAG_X4_V3, "%s-wdn" % TAG_X4_V3)
|
|
|
|
model_path = [model_path, wdn_model_path]
|
|
|
|
dni_weight = [params.denoise, 1 - params.denoise]
|
|
|
|
|
|
|
|
logger.debug("loading Real ESRGAN upscale model from %s", model_path)
|
|
|
|
|
2023-12-27 14:47:06 +00:00
|
|
|
upsampler = RealESRGANWrapper(
|
2023-07-01 12:10:53 +00:00
|
|
|
scale=params.scale,
|
2023-12-27 15:01:38 +00:00
|
|
|
model_path=None,
|
2023-07-01 12:10:53 +00:00
|
|
|
dni_weight=dni_weight,
|
|
|
|
model=model,
|
|
|
|
tile=tile,
|
|
|
|
tile_pad=params.tile_pad,
|
|
|
|
pre_pad=params.pre_pad,
|
2023-12-27 14:47:06 +00:00
|
|
|
half=("torch-fp16" in server.optimizations),
|
2023-12-27 15:01:38 +00:00
|
|
|
device=device.torch_str(),
|
2023-07-01 12:10:53 +00:00
|
|
|
)
|
|
|
|
|
2023-07-03 16:33:56 +00:00
|
|
|
server.cache.set(ModelTypes.upscaling, cache_key, upsampler)
|
2023-07-01 12:10:53 +00:00
|
|
|
run_gc([device])
|
|
|
|
|
|
|
|
return upsampler
|
|
|
|
|
|
|
|
def run(
|
|
|
|
self,
|
2023-07-15 23:54:54 +00:00
|
|
|
worker: WorkerContext,
|
2023-07-01 12:10:53 +00:00
|
|
|
server: ServerContext,
|
|
|
|
stage: StageParams,
|
|
|
|
_params: ImageParams,
|
2023-11-19 00:08:38 +00:00
|
|
|
sources: StageResult,
|
2023-07-01 12:10:53 +00:00
|
|
|
*,
|
|
|
|
upscale: UpscaleParams,
|
|
|
|
stage_source: Optional[Image.Image] = None,
|
|
|
|
**kwargs,
|
2023-11-19 00:08:38 +00:00
|
|
|
) -> StageResult:
|
2023-07-01 12:10:53 +00:00
|
|
|
logger.info("upscaling image with Real ESRGAN: x%s", upscale.scale)
|
|
|
|
|
2023-11-19 00:08:38 +00:00
|
|
|
upsampler = self.load(
|
|
|
|
server, upscale, worker.get_device(), tile=stage.tile_size
|
|
|
|
)
|
2023-07-01 12:10:53 +00:00
|
|
|
|
2023-11-19 00:08:38 +00:00
|
|
|
outputs = []
|
2024-01-06 02:11:58 +00:00
|
|
|
for source in sources.as_arrays():
|
2023-11-19 00:08:38 +00:00
|
|
|
output, _ = upsampler.enhance(source, outscale=upscale.outscale)
|
|
|
|
logger.info("final output image size: %s", output.shape)
|
2023-07-04 18:29:58 +00:00
|
|
|
outputs.append(output)
|
|
|
|
|
2024-01-06 22:59:02 +00:00
|
|
|
for metadata in sources.metadata:
|
|
|
|
metadata.upscale = upscale
|
|
|
|
|
2024-01-04 15:01:19 +00:00
|
|
|
return StageResult(arrays=outputs, metadata=sources.metadata)
|