app_haircolor.py

import spaces
import gradio as gr
import time
import torch
import numpy as np
import cv2

from PIL import Image
from segment_utils import(
    segment_image,
    restore_result,
)
from gfpgan.utils import GFPGANer
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from realesrgan.utils import RealESRGANer


DEFAULT_SRC_PROMPT = "a woman"
DEFAULT_EDIT_PROMPT = "a woman, with blue hair, 8k, high quality"

DEFAULT_CATEGORY = "hair"

device = "cuda" if torch.cuda.is_available() else "cpu"

def create_demo() -> gr.Blocks:
    from inversion_run_realvxl import run as realvxl_run
    model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
    model_path = 'realesr-general-x4v3.pth'
    half = True if torch.cuda.is_available() else False
    upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)

    @spaces.GPU(duration=10)
    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,
        lineart_scale: float,
        canny_scale: float,
    ):
        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)
        run_model = realvxl_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 = enhance(res_image)
        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


    def enhance(
        pil_image: Image,
    ):
        img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)

        h, w = img.shape[0:2]
        if h < 300:
            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
        output, _ = upsampler.enhance(img, outscale=2)
        pil_output = Image.fromarray(cv2.cvtColor(output, cv2.COLOR_BGR2RGB))

        return pil_output

    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.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")
                with gr.Accordion("Advanced Options", open=False):
                    guidance_scale = gr.Slider(minimum=0, maximum=20, value=1, step=0.5, label="Guidance Scale")
                    generate_size = gr.Number(label="Generate Size", value=512)
                    mask_expansion = gr.Number(label="Mask Expansion", value=10, visible=True)
                    mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation")
                    lineart_scale = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.1, label="Lineart Weights", visible=True)
                    canny_scale = gr.Slider(minimum=0, maximum=1, value=0.5, step=0.1, label="Canny Weights", visible=True)
            with gr.Column():
                seed = gr.Number(label="Seed", value=8)
                w1 = gr.Number(label="W1", value=2)
                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", 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", type="pil", interactive=False)
                enhanced_image = gr.Image(label="Enhanced Image", 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", 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, lineart_scale, canny_scale],
            outputs=[enhanced_image, generated_image, generated_cost],
        ).success(
            fn=restore_result,
            inputs=[croper, category, enhanced_image],
            outputs=[restored_image, download_path],
        )

    return demo