clean up draft CodeFormer stage
This commit is contained in:
parent
54dd34d211
commit
35681efc1b
|
@ -8,7 +8,7 @@ from PIL import Image
|
|||
from torchvision.transforms.functional import normalize
|
||||
|
||||
from ..device_pool import JobContext
|
||||
from ..params import ImageParams, StageParams
|
||||
from ..params import ImageParams, StageParams, UpscaleParams
|
||||
from ..utils import ServerContext
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
@ -27,9 +27,15 @@ def correct_codeformer(
|
|||
stage: StageParams,
|
||||
params: ImageParams,
|
||||
source_image: Image.Image,
|
||||
*,
|
||||
upscale: UpscaleParams = None,
|
||||
**kwargs,
|
||||
) -> Image.Image:
|
||||
ARCH_REGISTRY = {}
|
||||
bg_upsampler = None
|
||||
face_upsampler = None
|
||||
model = "TODO"
|
||||
w = None
|
||||
|
||||
# ------------------ set up CodeFormer restorer -------------------
|
||||
net = ARCH_REGISTRY.get("CodeFormer")(
|
||||
|
@ -47,7 +53,8 @@ def correct_codeformer(
|
|||
progress=True,
|
||||
file_name=None,
|
||||
)
|
||||
checkpoint = torch.load(ckpt_path)["params_ema"]
|
||||
checkpoint = torch.load(ckpt_path)
|
||||
checkpoint = checkpoint["params_ema"]
|
||||
net.load_state_dict(checkpoint)
|
||||
net.eval()
|
||||
|
||||
|
@ -96,7 +103,7 @@ def correct_codeformer(
|
|||
# upsample the background
|
||||
if bg_upsampler is not None:
|
||||
# Now only support RealESRGAN for upsampling background
|
||||
bg_img = bg_upsampler.enhance(img, outscale=args.upscale)[0]
|
||||
bg_img = bg_upsampler.enhance(source_image, outscale=upscale.scale)[0]
|
||||
else:
|
||||
bg_img = None
|
||||
|
||||
|
|
|
@ -1,2 +1,7 @@
|
|||
[metadata]
|
||||
description-file = README.md
|
||||
|
||||
[flake8]
|
||||
ignore = E203, W503
|
||||
max-line-length = 160
|
||||
per-file-ignores = __init__.py:F401
|
|
@ -14,6 +14,10 @@
|
|||
"settings": {
|
||||
"cSpell.words": [
|
||||
"astype",
|
||||
"basicsr",
|
||||
"ckpt",
|
||||
"codebook",
|
||||
"codeformer",
|
||||
"CUDA",
|
||||
"ddim",
|
||||
"ddpm",
|
||||
|
@ -43,11 +47,13 @@
|
|||
"outpainting",
|
||||
"outscale",
|
||||
"pndm",
|
||||
"pretrain",
|
||||
"pretrained",
|
||||
"protobuf",
|
||||
"randn",
|
||||
"realesr",
|
||||
"resrgan",
|
||||
"retinaface",
|
||||
"rocm",
|
||||
"RRDB",
|
||||
"runwayml",
|
||||
|
@ -63,10 +69,14 @@
|
|||
"timestep",
|
||||
"timesteps",
|
||||
"tojson",
|
||||
"torchvision",
|
||||
"uncond",
|
||||
"unet",
|
||||
"unsqueeze",
|
||||
"untruncated",
|
||||
"upsample",
|
||||
"upsampler",
|
||||
"upsampling",
|
||||
"upscaling",
|
||||
"venv",
|
||||
"virtualenv",
|
||||
|
|
Loading…
Reference in New Issue