|
|
import spaces |
|
|
import gradio as gr |
|
|
import time |
|
|
import torch |
|
|
import os |
|
|
import gc |
|
|
|
|
|
from PIL import Image, ImageEnhance, ImageFilter |
|
|
from segment_utils import( |
|
|
segment_image, |
|
|
restore_result_and_save, |
|
|
) |
|
|
from enhance_utils import enhance_sd_image |
|
|
from inversion_run_base import run as base_run |
|
|
|
|
|
DEFAULT_SRC_PROMPT = "a person" |
|
|
DEFAULT_EDIT_PROMPT = "a person with perfect face" |
|
|
|
|
|
DEFAULT_CATEGORY = "face" |
|
|
|
|
|
filter_names = [ |
|
|
"NONE", |
|
|
"DETAIL", |
|
|
"SMOOTH", |
|
|
"SMOOTH_MORE", |
|
|
"SHARPEN", |
|
|
"EDGE_ENHANCE", |
|
|
"EDGE_ENHANCE_MORE", |
|
|
] |
|
|
|
|
|
@spaces.GPU(duration=10) |
|
|
@torch.inference_mode() |
|
|
@torch.no_grad() |
|
|
def image_to_image( |
|
|
input_image: Image, |
|
|
input_image_prompt: str, |
|
|
edit_prompt: str, |
|
|
seed: int, |
|
|
w1: float, |
|
|
num_steps: int, |
|
|
start_step: int, |
|
|
guidance_scale: float, |
|
|
brightness: float = 1.0, |
|
|
color: float = 1.0, |
|
|
contrast: float = 1.0, |
|
|
sharpness: float = 1.0, |
|
|
filter: str = "NONE", |
|
|
): |
|
|
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) |
|
|
target_area_image = input_image |
|
|
|
|
|
run_model = base_run |
|
|
try: |
|
|
res_image = run_model( |
|
|
target_area_image, |
|
|
input_image_prompt, |
|
|
edit_prompt , |
|
|
seed, |
|
|
w1, |
|
|
w2, |
|
|
num_steps, |
|
|
start_step, |
|
|
guidance_scale, |
|
|
) |
|
|
run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'run_sd_model done') |
|
|
|
|
|
finally: |
|
|
torch.cuda.empty_cache() |
|
|
run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'cuda_empty_cache done') |
|
|
|
|
|
enhanced_image = res_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') |
|
|
|
|
|
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') |
|
|
|
|
|
enhancer = ImageEnhance.Brightness(enhanced_image) |
|
|
enhanced_image = enhancer.enhance(brightness) |
|
|
enhancer = ImageEnhance.Color(enhanced_image) |
|
|
enhanced_image = enhancer.enhance(color) |
|
|
enhancer = ImageEnhance.Contrast(enhanced_image) |
|
|
enhanced_image = enhancer.enhance(contrast) |
|
|
enhancer = ImageEnhance.Sharpness(enhanced_image) |
|
|
enhanced_image = enhancer.enhance(sharpness) |
|
|
|
|
|
if filter == "NONE": |
|
|
pass |
|
|
elif filter == "DETAIL": |
|
|
enhanced_image = enhanced_image.filter(ImageFilter.DETAIL) |
|
|
elif filter == "SMOOTH": |
|
|
enhanced_image = enhanced_image.filter(ImageFilter.SMOOTH) |
|
|
elif filter == "SMOOTH_MORE": |
|
|
enhanced_image = enhanced_image.filter(ImageFilter.SMOOTH_MORE) |
|
|
elif filter == "SHARPEN": |
|
|
enhanced_image = enhanced_image.filter(ImageFilter.SHARPEN) |
|
|
elif filter == "EDGE_ENHANCE": |
|
|
enhanced_image = enhanced_image.filter(ImageFilter.EDGE_ENHANCE) |
|
|
elif filter == "EDGE_ENHANCE_MORE": |
|
|
enhanced_image = enhanced_image.filter(ImageFilter.EDGE_ENHANCE_MORE) |
|
|
|
|
|
run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str, 'image_enhance done') |
|
|
|
|
|
return enhanced_image, time_cost_str |
|
|
|
|
|
def get_time_cost( |
|
|
run_task_time, |
|
|
time_cost_str, |
|
|
step: str = '' |
|
|
): |
|
|
now_time = int(time.time()*1000) |
|
|
if run_task_time == 0: |
|
|
time_cost_str = 'start' |
|
|
else: |
|
|
if time_cost_str != '': |
|
|
time_cost_str += f'-->' |
|
|
time_cost_str += f'{now_time - run_task_time}' |
|
|
if step != '': |
|
|
time_cost_str += f'-->{step}' |
|
|
run_task_time = now_time |
|
|
return run_task_time, time_cost_str |
|
|
|
|
|
def resize_image(image, target_size = 1024): |
|
|
h, w = image.size |
|
|
if h >= w: |
|
|
w = int(w * target_size / h) |
|
|
h = target_size |
|
|
else: |
|
|
h = int(h * target_size / w) |
|
|
w = target_size |
|
|
return image.resize((w, h)) |
|
|
|
|
|
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) |
|
|
with gr.Row(): |
|
|
brightness = gr.Number(label="Brightness", value=1.0, minimum=0.0, maximum=2.0, step=0.01) |
|
|
color = gr.Number(label="Color", value=1.0, minimum=0.0, maximum=2.0, step=0.01) |
|
|
contrast = gr.Number(label="Contrast", value=1.0, minimum=0.0, maximum=2.0, step=0.01) |
|
|
sharpness = gr.Number(label="Sharpness", value=1.0, minimum=0.0, maximum=2.0, step=0.01) |
|
|
with gr.Accordion("Advanced Options", open=False): |
|
|
category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False) |
|
|
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") |
|
|
with gr.Column(): |
|
|
num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps") |
|
|
start_step = gr.Slider(minimum=1, maximum=100, value=15, step=1, label="Start Step") |
|
|
filter = gr.Dropdown(choices=filter_names, label="Filter", value="NONE") |
|
|
g_btn = gr.Button("Edit Image") |
|
|
with gr.Accordion("Advanced Options", open=False): |
|
|
guidance_scale = gr.Slider(minimum=0, maximum=20, value=0, step=0.5, label="Guidance Scale") |
|
|
seed = gr.Number(label="Seed", value=8) |
|
|
w1 = gr.Number(label="W1", value=1.5) |
|
|
generate_size = gr.Number(label="Generate Size", value=1024) |
|
|
|
|
|
with gr.Row(): |
|
|
with gr.Column(): |
|
|
origin_area_image = gr.Image(label="Origin Area Image", format="png", type="pil", interactive=False) |
|
|
input_image = gr.Image(label="Input Image", type="pil", interactive=True) |
|
|
with gr.Column(): |
|
|
enhanced_image = gr.Image(label="Enhanced Image", format="png", type="pil", interactive=False) |
|
|
restored_image = gr.Image(label="Restored Image", format="png", type="pil", interactive=False) |
|
|
download_path = gr.File(label="Download the output image", interactive=False) |
|
|
generated_cost = gr.Textbox(label="Time cost by step (ms):", visible=True, interactive=False) |
|
|
|
|
|
g_btn.click( |
|
|
fn=segment_image, |
|
|
inputs=[input_image, 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, brightness, color, contrast, sharpness, filter], |
|
|
outputs=[enhanced_image, generated_cost], |
|
|
).success( |
|
|
fn=restore_result_and_save, |
|
|
inputs=[cropper, category, enhanced_image, save_quality], |
|
|
outputs=[restored_image, download_path], |
|
|
) |
|
|
|
|
|
return demo |
