add pre enhance

app_base.py

@@ -10,8 +10,8 @@ from segment_utils import(
 )
 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"
+DEFAULT_SRC_PROMPT = "a person"
+DEFAULT_EDIT_PROMPT = "a person with perfect face"
 
 DEFAULT_CATEGORY = "face"
 
@@ -31,13 +31,17 @@ def create_demo() -> gr.Blocks:
         start_step: int,
         guidance_scale: float,
         generate_size: int,
-        enhance_scale: int,
-        enhance_face: bool = True,
+        pre_enhance: bool = True,
+        pre_enhance_scale: int = 2,
     ):
         w2 = 1.0
         run_task_time = 0
         time_cost_str = ''
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+        if pre_enhance:
+            input_image = enhance_image(input_image, enhance_face=True, scale=pre_enhance_scale)
+            input_image = input_image.resize((generate_size, generate_size))
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
         run_model = base_run
         res_image = run_model(
             input_image,
@@ -52,8 +56,7 @@ def create_demo() -> gr.Blocks:
             guidance_scale,
         )
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-        enhance_mode = 0 if enhance_face else 2
-        enhanced_image = enhance_image(res_image, enhance_scale, enhance_mode)
+        enhanced_image = enhance_image(res_image)
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
 
         return enhanced_image, res_image, time_cost_str
@@ -77,18 +80,18 @@ def create_demo() -> gr.Blocks:
                 edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT)
                 category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False)
             with gr.Column():
-                num_steps = gr.Slider(minimum=1, maximum=100, value=30, step=1, label="Num Steps")
-                start_step = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Start Step")
+                num_steps = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Num Steps")
+                start_step = gr.Slider(minimum=1, maximum=100, value=30, step=1, label="Start Step")
                 with gr.Accordion("Advanced Options", open=False):
                     guidance_scale = gr.Slider(minimum=0, maximum=20, value=0, step=0.5, label="Guidance Scale")
-                    generate_size = gr.Number(label="Generate Size", value=1024)
+                    generate_size = gr.Number(label="Generate Size", value=512)
                     mask_expansion = gr.Number(label="Mask Expansion", value=50, visible=True)
                     mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation")
-                    enhance_scale = gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Enhance Scale")
-                    enhance_face = gr.Checkbox(label="Enhance Face", value=False)
+                    pre_enhance = gr.Checkbox(label="Pre Enhance", value=True)
+                    pre_enhance_scale = gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Pre Enhance Scale")
             with gr.Column():
                 seed = gr.Number(label="Seed", value=8)
-                w1 = gr.Number(label="W1", value=2)
+                w1 = gr.Number(label="W1", value=1.5)
                 g_btn = gr.Button("Edit Image")
                 
         with gr.Row():
@@ -109,7 +112,7 @@ def create_demo() -> gr.Blocks:
             outputs=[origin_area_image, croper],
         ).success(
             fn=image_to_image,
-            inputs=[origin_area_image, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, generate_size, enhance_scale, enhance_face],
+            inputs=[origin_area_image, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, generate_size, pre_enhance, pre_enhance_scale],
             outputs=[enhanced_image, generated_image, generated_cost],
         ).success(
             fn=restore_result,

enhance_utils.py

@@ -38,55 +38,21 @@ upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, ti
 
 face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2)
 
-
 def enhance_image(
-    input_image: Image,
-    scale: int,
-    enhance_mode: int,
-    keep_size: bool = False,
+    pil_image: Image,
+    enhance_face: bool = False,
+    scale: int = 2,
 ):
-    only_face = enhance_mode == 1
-    enhance_face = enhance_mode != 2
-
-    if enhance_mode == 1:
-        face_enhancer.upscale = scale
-        face_enhancer.bg_upsampler = None
-    elif enhance_mode == 2:
-        pass
-    else:
-        face_enhancer.upscale = scale
-        face_enhancer.bg_upsampler = upsampler
-    
-    img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
+    face_enhancer.upscale = scale
+    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)
-    
-    max_size = 3480 / scale
-    if h > max_size:
-        w = int(w * max_size / h)
-        h = max_size
-    
-    if w > max_size:
-        h = int(h * max_size / w)
-        w = max_size
-    
-    if h != img.shape[0] or w != img.shape[1]:
-        img = cv2.resize(img, (w, h), interpolation=cv2.INTER_LANCZOS4)
-    
     if enhance_face:
-        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=only_face, paste_back=True)
+        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=True, paste_back=True)
     else:
-        output, _ = upsampler.enhance(img, outscale=scale)
-    
-    h, w = img.shape[0:2]
-    interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
-    if keep_size:
-        output = cv2.resize(output, (w, h), interpolation=interpolation)
-    elif scale != 2:
-        output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
-    
+        output, _ = upsampler.enhance(img, outscale=2)
     pil_output = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))
 
     return pil_output

inversion_run_base.py

@@ -11,6 +11,8 @@ from config import get_config, get_num_steps_actual
 from functools import partial
 from compel import Compel, ReturnedEmbeddingsType
 
+from hidiffusion import apply_hidiffusion, remove_hidiffusion
+
 class Object(object):
     pass
 
@@ -55,6 +57,8 @@ pipeline.scheduler = DDPMScheduler.from_pretrained(
     subfolder="scheduler",
 )
 
+apply_hidiffusion(pipeline)
+
 config = get_config(args)
 
 compel_proc = Compel(

requirements.txt

@@ -14,4 +14,5 @@ git+https://github.com/XPixelGroup/BasicSR@master
 facexlib
 realesrgan
 controlnet_aux
-peft
+peft
+hidiffusion