Update app_2.py

app_2.py

@@ -10,16 +10,15 @@ import cv2
 import torch
 import torch.nn.functional as F
 import gradio as gr
-from itertools import chain
 
 from torchvision.transforms.functional import normalize
 
 from basicsr.utils import imwrite, img2tensor, tensor2img
 from basicsr.utils.download_util import load_file_from_url
 from facelib.utils.face_restoration_helper import FaceRestoreHelper
-from facelib.utils.misc import is_gray
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from basicsr.utils.realesrgan_utils import RealESRGANer
+from facelib.utils.misc import is_gray
 
 from basicsr.utils.registry import ARCH_REGISTRY
 
@@ -58,6 +57,9 @@ torch.hub.download_url_to_file(
 torch.hub.download_url_to_file(
     'https://replicate.com/api/models/sczhou/codeformer/files/7cf19c2c-e0cf-4712-9af8-cf5bdbb8d0ee/012.jpg',
     '05.jpg')
+torch.hub.download_url_to_file(
+    'https://raw.githubusercontent.com/sczhou/CodeFormer/master/inputs/cropped_faces/0729.png',
+    '06.png')
 
 def imread(img_path):
     img = cv2.imread(img_path)
@@ -102,19 +104,22 @@ codeformer_net.eval()
 
 os.makedirs('output', exist_ok=True)
 
-def inference(image, background_enhance, face_upsample, upscale, codeformer_fidelity):
+def inference(image, face_align, background_enhance, face_upsample, upscale, codeformer_fidelity):
     """Run a single prediction on the model"""
     try: # global try
         # take the default setting for the demo
-        has_aligned = False
         only_center_face = False
         draw_box = False
         detection_model = "retinaface_resnet50"
+
         print('Inp:', image, background_enhance, face_upsample, upscale, codeformer_fidelity)
-        
-        if background_enhance is None: background_enhance = True
-        if face_upsample is None: face_upsample = True
-        if upscale is None: upscale = 2
+        face_align = face_align if face_align is not None else True
+        background_enhance = background_enhance if background_enhance is not None else True
+        face_upsample = face_upsample if face_upsample is not None else True
+        upscale = upscale if (upscale is not None and upscale > 0) else 2
+
+        has_aligned = not face_align
+        upscale = 1 if has_aligned else upscale
 
         img = cv2.imread(str(image), cv2.IMREAD_COLOR)
         print('\timage size:', img.shape)
@@ -161,9 +166,7 @@ def inference(image, background_enhance, face_upsample, upscale, codeformer_fide
         # face restoration for each cropped face
         for idx, cropped_face in enumerate(face_helper.cropped_faces):
             # prepare data
-            cropped_face_t = img2tensor(
-                cropped_face / 255.0, bgr2rgb=True, float32=True
-            )
+            cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
             normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
             cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
 
@@ -177,12 +180,10 @@ def inference(image, background_enhance, face_upsample, upscale, codeformer_fide
                 torch.cuda.empty_cache()
             except RuntimeError as error:
                 print(f"Failed inference for CodeFormer: {error}")
-                restored_face = tensor2img(
-                    cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
-                )
+                restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
 
             restored_face = restored_face.astype("uint8")
-            face_helper.add_restored_face(restored_face)
+            face_helper.add_restored_face(restored_face, cropped_face)
 
         # paste_back
         if not has_aligned:
@@ -204,6 +205,8 @@ def inference(image, background_enhance, face_upsample, upscale, codeformer_fide
                 restored_img = face_helper.paste_faces_to_input_image(
                     upsample_img=bg_img, draw_box=draw_box
                 )
+        else:
+            restored_img = restored_face
 
         # save restored img
         save_path = f'output/out.png'
@@ -257,6 +260,12 @@ If you have any questions, please feel free to reach me out at <b>shangchenzhou@
 td {
     padding-right: 0px !important;
 }
+
+.gradio-container-4-37-2 .prose table, .gradio-container-4-37-2 .prose tr, .gradio-container-4-37-2 .prose td, .gradio-container-4-37-2 .prose th {
+    border: 0px solid #ffffff;
+    border-bottom: 0px solid #ffffff;
+}
+
 </style>
 
 <table>
@@ -269,38 +278,31 @@ td {
 <center><img src='https://api.infinitescript.com/badgen/count?name=sczhou/CodeFormer&ltext=Visitors&color=6dc9aa' alt='visitors'></center>
 """
 
-with gr.Blocks() as demo:
-    gr.Markdown(title)
-    gr.Markdown(description)
-    with gr.Box():
-        with gr.Column():
-            input_img = gr.Image(type="filepath", label="Input")
-            background_enhance = gr.Checkbox(value=True, label="Background_Enhance")
-            face_enhance = gr.Checkbox(value=True, label="Face_Upsample")
-            upscale_factor = gr.Number(value=2, label="Rescaling_Factor (up to 4)")
-            codeformer_fidelity = gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity (0 for better quality, 1 for better identity)')
-            submit = gr.Button('Enhance Image')
-        with gr.Column():
-            output_img = gr.Image(type="numpy", label="Output").style(height='auto')
-            
-    inps = [input_img, background_enhance, face_enhance, upscale_factor, codeformer_fidelity]
-    submit.click(fn=inference, inputs=inps, outputs=[output_img])
-            
-    ex = gr.Examples([
-        ['01.png', True, True, 2, 0.7],
-        ['02.jpg', True, True, 2, 0.7],
-        ['03.jpg', True, True, 2, 0.7],
-        ['04.jpg', True, True, 2, 0.1],
-        ['05.jpg', True, True, 2, 0.1]
+demo = gr.Interface(
+    inference, [
+        gr.Image(type="filepath", label="Input"),
+        gr.Checkbox(value=True, label="Pre_Face_Align"),
+        gr.Checkbox(value=True, label="Background_Enhance"),
+        gr.Checkbox(value=True, label="Face_Upsample"),
+        gr.Number(value=2, label="Rescaling_Factor (up to 4)"),
+        gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity (0 for better quality, 1 for better identity)')
+    ], [
+        gr.Image(type="numpy", label="Output")
+    ],
+    title=title,
+    description=description,
+    article=article,       
+    examples=[
+        ['01.png', True, True, True, 2, 0.7],
+        ['02.jpg', True, True, True, 2, 0.7],
+        ['03.jpg', True, True, True, 2, 0.7],
+        ['04.jpg', True, True, True, 2, 0.1],
+        ['05.jpg', True, True, True, 2, 0.1],
+        ['06.png', False, True, True, 1, 0.5]
       ],
-        fn=inference,
-        inputs=inps,
-        outputs=[output_img],
-        cache_examples=True)
-    
-    gr.Markdown(article)
-    
-    
+    concurrency_limit=2
+    )
+
 DEBUG = os.getenv('DEBUG') == '1'
-demo.queue(api_open=False, concurrency_count=2, max_size=10)
-demo.launch(debug=DEBUG)
+# demo.launch(debug=DEBUG)
+demo.launch(debug=DEBUG, share=True)