app_base.py

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