From 35681efc1b428dad4d5813bc7899e9338e51cd63 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sun, 5 Feb 2023 08:06:50 -0600
Subject: [PATCH] clean up draft CodeFormer stage

---
 api/onnx_web/chain/correct_codeformer.py | 13 ++++++++++---
 api/setup.cfg                            |  7 ++++++-
 onnx-web.code-workspace                  | 10 ++++++++++
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/api/onnx_web/chain/correct_codeformer.py b/api/onnx_web/chain/correct_codeformer.py
index 7104255f..34fa34d9 100644
--- a/api/onnx_web/chain/correct_codeformer.py
+++ b/api/onnx_web/chain/correct_codeformer.py
@@ -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
 
diff --git a/api/setup.cfg b/api/setup.cfg
index 224a7795..053f652f 100644
--- a/api/setup.cfg
+++ b/api/setup.cfg
@@ -1,2 +1,7 @@
 [metadata]
-description-file = README.md
\ No newline at end of file
+description-file = README.md
+
+[flake8]
+ignore = E203, W503
+max-line-length = 160
+per-file-ignores = __init__.py:F401
\ No newline at end of file
diff --git a/onnx-web.code-workspace b/onnx-web.code-workspace
index 42dd5d87..20d949ba 100644
--- a/onnx-web.code-workspace
+++ b/onnx-web.code-workspace
@@ -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",