import spaces import gradio as gr import time import torch import tempfile import os import gc from loading_utils import load_image from segment_utils import( segment_image, restore_result, ) 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" def image_to_image( 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, ): 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') 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') 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') return output_path, restored_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 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: 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.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") 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") 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(): input_image_path = gr.Image(label="Input Image", type="filepath", interactive=True) 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) 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], ) return demo