xl_fb / app_base.py
zhiweili
add upscale
b38f27e
raw
history blame
6.15 kB
import spaces
import gradio as gr
import time
import torch
from PIL import Image
from segment_utils import(
segment_image,
restore_result,
)
from upscale import upscale_image
DEFAULT_SRC_PROMPT = "a person"
DEFAULT_EDIT_PROMPT = "a person with perfect face"
DEFAULT_CATEGORY = "face"
device = "cuda" if torch.cuda.is_available() else "cpu"
def create_demo() -> gr.Blocks:
from inversion_run_base import run as base_run
@spaces.GPU(duration=30)
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,
generate_size: int,
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_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
res_image = run_model(
input_image,
input_image_prompt,
edit_prompt,
generate_size,
seed,
w1,
w2,
num_steps,
start_step,
guidance_scale,
)
run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
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
def get_time_cost(run_task_time, time_cost_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}'
run_task_time = now_time
return run_task_time, time_cost_str
with gr.Blocks() as demo:
croper = 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)
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=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")
with gr.Column():
seed = gr.Number(label="Seed", value=8)
w1 = gr.Number(label="W1", value=1.5)
g_btn = gr.Button("Edit Image")
with gr.Row():
with gr.Column():
input_image = gr.Image(label="Input Image", type="pil")
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)
with gr.Column():
origin_area_image = gr.Image(label="Origin Area Image", format="png", type="pil", 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)
generated_image = gr.Image(label="Generated Image", format="png", type="pil", interactive=False)
g_btn.click(
fn=segment_image,
inputs=[input_image, category, generate_size, mask_expansion, mask_dilation],
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, 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,
inputs=[croper, category, enhanced_image],
outputs=[restored_image, download_path],
)
return demo