add upscale

app_base.py

@@ -8,7 +8,7 @@ from segment_utils import(
     segment_image,
     restore_result,
 )
-from enhance_utils import enhance_image
+from upscale import upscale_image
 
 DEFAULT_SRC_PROMPT = "a person"
 DEFAULT_EDIT_PROMPT = "a person with perfect face"
@@ -31,15 +31,24 @@ def create_demo() -> gr.Blocks:
         start_step: int,
         guidance_scale: float,
         generate_size: int,
-        pre_enhance: bool = True,
-        pre_enhance_scale: int = 2,
+        upscale_prompt: str,
+        upscale_start_size: int = 256,
+        upscale_steps: int = 10,
+        pre_upscale: bool = True,
+        pre_upscale_start_size: int = 128,
+        pre_upscale_steps: int = 30,
     ):
         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)
+        if pre_upscale:
+            input_image = upscale_image(
+                input_image,
+                upscale_prompt,
+                start_size=pre_upscale_start_size,
+                upscale_steps=pre_upscale_steps,
+            )
             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
@@ -56,7 +65,12 @@ def create_demo() -> gr.Blocks:
             guidance_scale,
         )
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-        enhanced_image = enhance_image(res_image)
+        enhanced_image = upscale_image(
+            res_image,
+            upscale_prompt,
+            start_size=upscale_start_size,
+            upscale_steps=upscale_steps,
+        )
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
 
         return enhanced_image, res_image, time_cost_str
@@ -79,16 +93,21 @@ def create_demo() -> gr.Blocks:
                 input_image_prompt = gr.Textbox(lines=1, label="Input Image Prompt", value=DEFAULT_SRC_PROMPT)
                 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.Accordion("Advanced Options", open=False):
+                    upscale_prompt = gr.Textbox(lines=1, label="Upscale Prompt", value="a person with pefect face")
+                    upscale_start_size = gr.Number(label="Upscale Start Size", value=256)
+                    upscale_steps = gr.Number(label="Upscale Steps", value=10)
+                    pre_upscale = gr.Checkbox(label="Pre Upscale", value=True)
+                    pre_upscale_start_size = gr.Number(label="Pre Upscale Start Size", value=128)
+                    pre_upscale_steps = gr.Number(label="Pre Upscale Steps", value=30)
             with gr.Column():
                 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=512)
+                    generate_size = gr.Number(label="Generate Size", value=256)
                     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")
-                    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=1.5)
@@ -112,7 +131,9 @@ 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, pre_enhance, pre_enhance_scale],
+            inputs=[origin_area_image, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step,
+                guidance_scale, generate_size, upscale_prompt, upscale_start_size, upscale_steps,
+                pre_upscale, pre_upscale_start_size, pre_upscale_steps],
             outputs=[enhanced_image, generated_image, generated_cost],
         ).success(
             fn=restore_result,

upscale.py

@@ -0,0 +1,25 @@
+import torch
+
+from PIL import Image
+from diffusers import StableDiffusionUpscalePipeline
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_id = "stabilityai/stable-diffusion-x4-upscaler"
+upscale_pipe = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+upscale_pipe = upscale_pipe.to(device)
+
+def upscale_image(
+    input_image: Image,
+    prompt: str,
+    start_size: int = 128,
+    upscale_steps: int = 30,
+):
+    input_image = input_image.resize((start_size, start_size))
+    upscaled_image = upscale_pipe(
+        prompt=prompt, 
+        image=input_image,
+        num_inference_steps=upscale_steps,
+    ).images[0]
+    
+    return upscaled_image