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-03-11 00:15:31 +00:00
|
|
|
from typing import Optional
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-01-31 14:16:57 +00:00
|
|
|
import numpy as np
|
2023-02-05 13:53:26 +00:00
|
|
|
from PIL import Image
|
|
|
|
|
2023-02-06 14:07:06 +00:00
|
|
|
from ..params import DeviceParams, ImageParams, StageParams, UpscaleParams
|
2023-02-26 20:15:30 +00:00
|
|
|
from ..server import ServerContext
|
2023-02-19 02:28:21 +00:00
|
|
|
from ..utils import run_gc
|
2023-02-26 05:49:39 +00:00
|
|
|
from ..worker import WorkerContext
|
2023-01-31 14:16:57 +00:00
|
|
|
|
2023-01-28 23:09:19 +00:00
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-02-06 23:26:51 +00:00
|
|
|
def load_gfpgan(
|
|
|
|
server: ServerContext,
|
2023-02-14 00:10:11 +00:00
|
|
|
_stage: StageParams,
|
2023-02-06 23:26:51 +00:00
|
|
|
upscale: UpscaleParams,
|
2023-02-17 00:11:35 +00:00
|
|
|
device: DeviceParams,
|
2023-02-06 23:26:51 +00:00
|
|
|
):
|
2023-02-14 04:23:17 +00:00
|
|
|
# must be within the load function for patch to take effect
|
2023-04-11 13:26:21 +00:00
|
|
|
# TODO: rewrite and remove
|
2023-02-14 04:23:17 +00:00
|
|
|
from gfpgan import GFPGANer
|
|
|
|
|
|
|
|
face_path = path.join(server.cache_path, "%s.pth" % (upscale.correction_model))
|
2023-02-14 00:10:11 +00:00
|
|
|
cache_key = (face_path,)
|
|
|
|
cache_pipe = server.cache.get("gfpgan", cache_key)
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-02-14 00:10:11 +00:00
|
|
|
if cache_pipe is not None:
|
2023-02-05 13:53:26 +00:00
|
|
|
logger.info("reusing existing GFPGAN pipeline")
|
2023-02-14 00:10:11 +00:00
|
|
|
return cache_pipe
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-02-05 13:53:26 +00:00
|
|
|
logger.debug("loading GFPGAN model from %s", face_path)
|
2023-02-03 05:34:02 +00:00
|
|
|
|
2023-01-31 23:08:30 +00:00
|
|
|
# TODO: find a way to pass the ONNX model to underlying architectures
|
2023-01-28 05:28:14 +00:00
|
|
|
gfpgan = GFPGANer(
|
2023-02-05 13:53:26 +00:00
|
|
|
arch="clean",
|
2023-02-12 19:50:28 +00:00
|
|
|
bg_upsampler=None,
|
2023-01-28 05:28:14 +00:00
|
|
|
channel_multiplier=2,
|
2023-04-11 13:34:21 +00:00
|
|
|
device=device.torch_str(),
|
2023-02-06 23:13:37 +00:00
|
|
|
model_path=face_path,
|
|
|
|
upscale=upscale.face_outscale,
|
2023-02-05 13:53:26 +00:00
|
|
|
)
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-02-14 00:10:11 +00:00
|
|
|
server.cache.set("gfpgan", cache_key, gfpgan)
|
2023-02-17 00:11:35 +00:00
|
|
|
run_gc([device])
|
2023-01-28 05:28:14 +00:00
|
|
|
|
|
|
|
return gfpgan
|
|
|
|
|
|
|
|
|
|
|
|
def correct_gfpgan(
|
2023-02-26 05:49:39 +00:00
|
|
|
job: WorkerContext,
|
2023-02-05 03:17:39 +00:00
|
|
|
server: ServerContext,
|
2023-02-06 23:13:37 +00:00
|
|
|
stage: StageParams,
|
2023-01-28 05:28:14 +00:00
|
|
|
_params: ImageParams,
|
2023-02-18 22:35:57 +00:00
|
|
|
source: Image.Image,
|
2023-01-28 05:28:14 +00:00
|
|
|
*,
|
|
|
|
upscale: UpscaleParams,
|
2023-03-11 00:15:31 +00:00
|
|
|
stage_source: Optional[Image.Image] = None,
|
2023-01-29 04:31:34 +00:00
|
|
|
**kwargs,
|
2023-01-28 18:42:02 +00:00
|
|
|
) -> Image.Image:
|
2023-02-18 22:27:48 +00:00
|
|
|
upscale = upscale.with_args(**kwargs)
|
2023-02-19 04:11:44 +00:00
|
|
|
source = stage_source or source
|
2023-02-18 22:27:48 +00:00
|
|
|
|
2023-01-28 05:28:14 +00:00
|
|
|
if upscale.correction_model is None:
|
2023-02-05 13:53:26 +00:00
|
|
|
logger.warn("no face model given, skipping")
|
2023-02-18 22:35:57 +00:00
|
|
|
return source
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-02-05 13:53:26 +00:00
|
|
|
logger.info("correcting faces with GFPGAN model: %s", upscale.correction_model)
|
2023-02-06 14:07:06 +00:00
|
|
|
device = job.get_device()
|
2023-02-06 23:13:37 +00:00
|
|
|
gfpgan = load_gfpgan(server, stage, upscale, device)
|
2023-01-28 05:28:14 +00:00
|
|
|
|
2023-02-18 22:35:57 +00:00
|
|
|
output = np.array(source)
|
2023-01-28 05:28:14 +00:00
|
|
|
_, _, output = gfpgan.enhance(
|
2023-02-05 13:53:26 +00:00
|
|
|
output,
|
|
|
|
has_aligned=False,
|
|
|
|
only_center_face=False,
|
|
|
|
paste_back=True,
|
|
|
|
weight=upscale.face_strength,
|
|
|
|
)
|
|
|
|
output = Image.fromarray(output, "RGB")
|
2023-01-28 05:28:14 +00:00
|
|
|
|
|
|
|
return output
|