add brightness

app_base.py

@@ -2,15 +2,13 @@ import spaces
 import gradio as gr
 import time
 import torch
-import tempfile
 import os
 import gc
 
-from loading_utils import load_image
-
+from PIL import Image, ImageEnhance
 from segment_utils import(
     segment_image,
-    restore_result,
+    restore_result_and_save,
 )
 from enhance_utils import enhance_sd_image
 from inversion_run_base import run as base_run
@@ -20,8 +18,11 @@ DEFAULT_EDIT_PROMPT = "a person with perfect face"
 
 DEFAULT_CATEGORY = "face"
 
+@spaces.GPU(duration=10)
+@torch.inference_mode()
+@torch.no_grad()
 def image_to_image(
-    input_image_path: str,
+    input_image: Image,
     input_image_prompt: str,
     edit_prompt: str,
     seed: int,
@@ -29,35 +30,14 @@ def image_to_image(
     num_steps: int,
     start_step: int,
     guidance_scale: float,
-    generate_size: int,
-    mask_expansion: int = 50,
-    mask_dilation: int = 2,
-    save_quality: int = 95,
-    enable_segment: bool = True,
+    brightness: float = 1.0,
 ):
-    segment_category = "face"
     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)
-    input_image = load_image(input_image_path)
-    icc_profile = input_image.info.get('icc_profile')
-    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'load_image done')
-
-    if enable_segment:
-        target_area_image, croper = segment_image(
-            input_image,
-            segment_category,
-            generate_size,
-            mask_expansion,
-            mask_dilation,
-        )
-    else:
-        target_area_image = resize_image(input_image, generate_size)
-        croper = None
-    
-    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'segment_image done')
+    target_area_image = input_image
     
     run_model = base_run
     try:
@@ -82,30 +62,16 @@ def image_to_image(
     enhanced_image = enhance_sd_image(res_image)
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'enhance_image done')
 
-    if enable_segment:
-        restored_image = restore_result(croper, segment_category, enhanced_image)
-    else:
-        restored_image = enhanced_image.resize(input_image.size)
-    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'restore_result done')
-
     torch.cuda.empty_cache()
     run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'cuda_empty_cache done')
     if os.getenv('ENABLE_GC', False):
         gc.collect()
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'gc_collect done')
 
-    extension = 'png'
-    if restored_image.mode == 'RGBA':
-        extension = 'png'
-    else:
-        extension = 'webp'
-
-    output_path = tempfile.mktemp(suffix=f".{extension}")
-    restored_image.save(output_path, format=extension, quality=save_quality, icc_profile=icc_profile)
-
-    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'save_image done')
+    enhancer = ImageEnhance.Brightness(enhanced_image)
+    enhanced_image = enhancer.enhance(brightness)
 
-    return output_path, restored_image, time_cost_str
+    return enhanced_image, time_cost_str
 
 def get_time_cost(
     run_task_time, 
@@ -134,49 +100,16 @@ def resize_image(image, target_size = 1024):
         w = target_size
     return image.resize((w, h))
 
-
-def infer(
-    input_image_path: str,
-    input_image_prompt: str,
-    edit_prompt: str,
-    seed: int,
-    w1: float,
-    num_steps: int,
-    start_step: int,
-    guidance_scale: float,
-    generate_size: int,
-    mask_expansion: int = 50,
-    mask_dilation: int = 2,
-    save_quality: int = 95,
-    enable_segment: bool = True,
-):
-    return image_to_image(
-        input_image_path,
-        input_image_prompt,
-        edit_prompt,
-        seed,
-        w1,
-        num_steps,
-        start_step,
-        guidance_scale,
-        generate_size,
-        mask_expansion,
-        mask_dilation,
-        save_quality,
-        enable_segment
-    )
-
-infer = spaces.GPU(infer)
-
 def create_demo() -> gr.Blocks:
 
     with gr.Blocks() as demo:
+        cropper = gr.State()
         with gr.Row():
             with gr.Column():
                 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)
+                brightness = gr.Slider(minimum=0, maximum=2, value=1.0, step=0.1, label="Brightness")
                 with gr.Accordion("Advanced Options", open=False):
-                    enable_segment = gr.Checkbox(label="Enable Segment", value=True)
                     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")
                     save_quality = gr.Slider(minimum=1, maximum=100, value=95, step=1, label="Save Quality")
@@ -192,18 +125,26 @@ def create_demo() -> gr.Blocks:
         
         with gr.Row():
             with gr.Column():
-                input_image_path = gr.Image(label="Input Image", type="filepath", interactive=True)
+                input_image = gr.Image(label="Input Image", type="pil", interactive=True)
+                origin_area_image = gr.Image(label="Origin Area Image", format="png", type="pil", interactive=False)
             with gr.Column():
                 restored_image = gr.Image(label="Restored Image", format="png", type="pil", interactive=False)
                 download_path = gr.File(label="Download the output image", interactive=False)
+                enhanced_image = gr.Image(label="Enhanced Image", format="png", type="pil", interactive=False)
                 generated_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, interactive=False)
-        
+
         g_btn.click(
-            fn=infer,
-            inputs=[input_image_path, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale, generate_size, mask_expansion, mask_dilation, save_quality, enable_segment],
-            outputs=[download_path, restored_image, generated_cost],
+            fn=segment_image,
+            inputs=[input_image, DEFAULT_CATEGORY, generate_size, mask_expansion, mask_dilation],
+            outputs=[origin_area_image, cropper],
+        ).success(
+            fn=image_to_image,
+            inputs=[origin_area_image, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step, guidance_scale],
+            outputs=[enhanced_image, generated_cost],
+        ).success(
+            fn=restore_result_and_save,
+            inputs=[cropper, DEFAULT_CATEGORY, enhanced_image, save_quality],
+            outputs=[restored_image, download_path],
         )
-    
-    
 
     return demo

segment_utils.py

@@ -1,6 +1,6 @@
 import numpy as np
 import mediapipe as mp
-import uuid
+import tempfile
 
 from PIL import Image
 from scipy.ndimage import binary_dilation
@@ -22,6 +22,28 @@ def restore_result(croper, category, generated_image):
 
     return restored_image
 
+def restore_result_and_save(croper, category, generated_image,save_quality=95):
+    square_length = croper.square_length
+    generated_image = generated_image.resize((square_length, square_length))
+
+    cropped_generated_image = generated_image.crop((croper.square_start_x, croper.square_start_y, croper.square_end_x, croper.square_end_y))
+    cropped_square_mask_image = get_restore_mask_image(croper, category, cropped_generated_image)
+
+    restored_image = croper.input_image.copy()
+    restored_image.paste(cropped_generated_image, (croper.origin_start_x, croper.origin_start_y), cropped_square_mask_image)
+
+    extension = 'png'
+    if restored_image.mode == 'RGBA':
+        extension = 'png'
+    else:
+        extension = 'webp'
+
+    icc_profile = croper.input_image.info.get('icc_profile')
+    output_path = tempfile.mktemp(suffix=f".{extension}")
+    restored_image.save(output_path, format=extension, quality=save_quality, icc_profile=icc_profile)
+
+    return restored_image, output_path
+
 def segment_image(input_image, category, input_size, mask_expansion, mask_dilation):
     mask_size = int(input_size)
     mask_expansion = int(mask_expansion)