add enhance utils

app.py

@@ -1,12 +1,12 @@
 import gradio as gr
 
-# from app_base import create_demo as create_demo_face
+from app_base import create_demo as create_demo_face
 from app_haircolor import create_demo as create_demo_haircolor
 
 with gr.Blocks(css="style.css") as demo:
     with gr.Tabs():
-        # with gr.Tab(label="Face"):
-        #     create_demo_face()
+        with gr.Tab(label="Face"):
+            create_demo_face()
         with gr.Tab(label="Hair Color"):
             create_demo_haircolor()
 

app_base.py

@@ -2,19 +2,13 @@ import spaces
 import gradio as gr
 import time
 import torch
-import os
-import numpy as np
-import cv2
 
 from PIL import Image
 from segment_utils import(
     segment_image,
     restore_result,
 )
-from gfpgan.utils import GFPGANer
-from basicsr.archs.srvgg_arch import SRVGGNetCompact
-from realesrgan.utils import RealESRGANer
-
+from enhance_utils import enhance_image
 
 DEFAULT_SRC_PROMPT = "a woman, photo"
 DEFAULT_EDIT_PROMPT = "a beautiful woman, photo, hollywood style face, 8k, high quality"
@@ -25,12 +19,6 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def create_demo() -> gr.Blocks:
     from inversion_run_base import run as base_run
-    model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
-    model_path = 'realesr-general-x4v3.pth'
-    half = True if torch.cuda.is_available() else False
-    upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
-
-    face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=1, arch='clean', channel_multiplier=2)
 
     @spaces.GPU(duration=10)
     def image_to_image(
@@ -65,7 +53,7 @@ def create_demo() -> gr.Blocks:
             adapter_weights,
         )
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-        enhanced_image = enhance(res_image, enhance_face)
+        enhanced_image = enhance_image(res_image, enhance_face)
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
 
         return enhanced_image, res_image, time_cost_str
@@ -81,24 +69,6 @@ def create_demo() -> gr.Blocks:
         run_task_time = now_time
         return run_task_time, time_cost_str
 
-
-    def enhance(
-        pil_image: Image,
-        enhance_face: bool = True,
-    ):
-        img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
-
-        h, w = img.shape[0:2]
-        if h < 300:
-            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
-        if enhance_face:
-            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=True, paste_back=True)
-        else:
-            output, _ = upsampler.enhance(img, outscale=2)
-        pil_output = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
-
-        return pil_output
-
     with gr.Blocks() as demo:
         croper = gr.State()
         with gr.Row():

app_haircolor.py

@@ -10,9 +10,7 @@ from segment_utils import(
     segment_image,
     restore_result,
 )
-from gfpgan.utils import GFPGANer
-from basicsr.archs.srvgg_arch import SRVGGNetCompact
-from realesrgan.utils import RealESRGANer
+from enhance_utils import enhance_image
 
 
 DEFAULT_SRC_PROMPT = "a woman"
@@ -23,11 +21,7 @@ DEFAULT_CATEGORY = "hair"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def create_demo() -> gr.Blocks:
-    from inversion_run_adapter import run as realvxl_run
-    model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
-    model_path = 'realesr-general-x4v3.pth'
-    half = True if torch.cuda.is_available() else False
-    upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+    from inversion_run_adapter import run as adapter_run
 
     @spaces.GPU(duration=10)
     def image_to_image(
@@ -48,7 +42,7 @@ def create_demo() -> gr.Blocks:
         run_task_time = 0
         time_cost_str = ''
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-        run_model = realvxl_run
+        run_model = adapter_run
         res_image = run_model(
             input_image,
             input_image_prompt,
@@ -65,7 +59,7 @@ def create_demo() -> gr.Blocks:
             sketch_scale,
         )
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-        enhanced_image = enhance(res_image)
+        enhanced_image = enhance_image(res_image, False)
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
 
         return enhanced_image, res_image, time_cost_str
@@ -81,20 +75,6 @@ def create_demo() -> gr.Blocks:
         run_task_time = now_time
         return run_task_time, time_cost_str
 
-
-    def enhance(
-        pil_image: Image,
-    ):
-        img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
-
-        h, w = img.shape[0:2]
-        if h < 300:
-            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
-        output, _ = upsampler.enhance(img, outscale=2)
-        pil_output = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
-
-        return pil_output
-
     with gr.Blocks() as demo:
         croper = gr.State()
         with gr.Row():

enhance_utils.py

@@ -0,0 +1,41 @@
+import os
+import torch
+import cv2
+import numpy as np
+
+from PIL import Image
+from gfpgan.utils import GFPGANer
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from realesrgan.utils import RealESRGANer
+
+os.system("pip freeze")
+if not os.path.exists('GFPGANv1.4.pth'):
+    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
+if not os.path.exists('realesr-general-x4v3.pth'):
+    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
+
+os.makedirs('output', exist_ok=True)
+
+model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+model_path = 'realesr-general-x4v3.pth'
+half = True if torch.cuda.is_available() else False
+upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+
+face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=1, arch='clean', channel_multiplier=2)
+
+def enhance_image(
+    pil_image: Image,
+    enhance_face: bool = True,
+):
+    img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+
+    h, w = img.shape[0:2]
+    if h < 300:
+        img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
+    if enhance_face:
+        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=True, paste_back=True)
+    else:
+        output, _ = upsampler.enhance(img, outscale=2)
+    pil_output = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
+
+    return pil_output

inversion_run_adapter.py

@@ -1,5 +1,4 @@
 import torch
-import os
 
 from diffusers import (
     DDPMScheduler,
@@ -20,14 +19,6 @@ from config import get_config, get_num_steps_actual
 from functools import partial
 from compel import Compel, ReturnedEmbeddingsType
 
-os.system("pip freeze")
-if not os.path.exists('GFPGANv1.4.pth'):
-    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
-if not os.path.exists('realesr-general-x4v3.pth'):
-    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
-
-os.makedirs('output', exist_ok=True)
-
 class Object(object):
     pass