add pre upscale

app_base.py

@@ -10,6 +10,8 @@ from segment_utils import(
 )
 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,16 +33,24 @@ def create_demo() -> gr.Blocks:
         start_step: int,
         guidance_scale: float,
         generate_size: int,
-        pre_enhance: bool = True,
-        pre_enhance_scale: int = 2,
+        enhance_scale: int = 2,
+        pre_upscale: bool = True,
+        upscale_prompt: str,
+        pre_upscale_steps: int = 10,
     ):
         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))
+        if pre_upscale:
+            pre_upscale_start_size = generate_size // 4
+            input_image = upscale_image(
+                input_image,
+                upscale_prompt,
+                start_size=pre_upscale_start_size,
+                upscale_steps=pre_upscale_steps,
+                seed=seed,
+            )
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
         run_model = base_run
         res_image = run_model(
@@ -56,7 +66,7 @@ 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 = enhance_image(res_image, scale = enhance_scale)
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
 
         return enhanced_image, res_image, time_cost_str
@@ -79,6 +89,11 @@ 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):
+                    enhance_scale = gr.Number(label="Enhance Scale", value=2)
+                    pre_upscale = gr.Checkbox(label="Pre Upscale", value=True)
+                    upscale_prompt = gr.Textbox(lines=1, label="Upscale Prompt", value="a person with pefect face")
+                    pre_upscale_steps = gr.Number(label="Pre Upscale Steps", value=10)
             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")
@@ -87,8 +102,6 @@ def create_demo() -> gr.Blocks:
                     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")
-                    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 +125,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, pre_enhance, pre_enhance_scale],
+            inputs=[origin_area_image, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, generate_size, enhance_scale, pre_upscale, upscale_prompt, pre_upscale_steps],
             outputs=[enhanced_image, generated_image, generated_cost],
         ).success(
             fn=restore_result,

upscale.py

@@ -0,0 +1,27 @@
+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,
+    seed: int = 42,
+):
+    generator = torch.Generator().manual_seed(seed)
+    input_image = input_image.resize((start_size, start_size))
+    upscaled_image = upscale_pipe(
+        prompt=prompt, 
+        image=input_image,
+        num_inference_steps=upscale_steps,
+        generator=generator,
+    ).images[0]
+    
+    return upscaled_image