app_haircolor_img2img.py

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

from PIL import Image
from segment_utils import(
    segment_image,
    restore_result,
)
from diffusers import (
    DiffusionPipeline,
    T2IAdapter,
    MultiAdapter,
    AutoencoderKL,
    EulerAncestralDiscreteScheduler,
)

from controlnet_aux import (
    CannyDetector,
    LineartDetector,
)

BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"

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

DEFAULT_EDIT_PROMPT = "RAW photo, Fujifilm XT3,  sharp hair, high resolution hair, hair tones, natural hair, magazine hair, white color hair"

DEFAULT_CATEGORY = "hair"

canny_detector = CannyDetector()
lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
lineart_detector = lineart_detector.to(DEVICE)


adapters = MultiAdapter(
    [
        T2IAdapter.from_pretrained(
            "TencentARC/t2i-adapter-lineart-sdxl-1.0",
            torch_dtype=torch.float16,
            varient="fp16",
        ),
        T2IAdapter.from_pretrained(
            "TencentARC/t2i-adapter-canny-sdxl-1.0",
            torch_dtype=torch.float16,
            varient="fp16",
        ),
    ]
)
adapters = adapters.to(torch.float16)

basepipeline = DiffusionPipeline.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16, 
    variant="fp16", 
    use_safetensors=True,
    vae=AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16),
    scheduler=EulerAncestralDiscreteScheduler.from_pretrained(BASE_MODEL, subfolder="scheduler"),
    adapter=adapters,
    custom_pipeline="./pipelines/pipeline_sdxl_adapter_img2img.py",
)

basepipeline = basepipeline.to(DEVICE)

basepipeline.enable_model_cpu_offload()

@spaces.GPU(duration=30)
def image_to_image(
    input_image: Image,
    edit_prompt: str,
    seed: int,
    num_steps: int,
    guidance_scale: float,
    strength: float,
    generate_size: int,
    cond_scale1: float = 1.2,
    cond_scale2: float = 1.2,
    lineart_detect:float = 0.375,
    canny_detect:float = 0.375,
):
    run_task_time = 0
    time_cost_str = ''
    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
    lineart_image = lineart_detector(input_image, int(generate_size * lineart_detect), generate_size)
    canny_image = canny_detector(input_image, int(generate_size * canny_detect), generate_size)

    cond_image = [lineart_image, canny_image]
    cond_scale = [cond_scale1, cond_scale2]

    generator = torch.Generator(device=DEVICE).manual_seed(seed)
    generated_image = basepipeline(
        generator=generator,
        prompt=edit_prompt,
        image=input_image,
        height=generate_size,
        width=generate_size,
        guidance_scale=guidance_scale,
        strength=strength,
        num_inference_steps=num_steps,
        adapter_image=cond_image,
        adapter_conditioning_scale=cond_scale,
    ).images[0]
    
    run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)

    return generated_image, time_cost_str

def make_inpaint_condition(image, image_mask):
    image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
    image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0

    assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
    image[image_mask > 0.5] = -1.0  # set as masked pixel
    image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
    image = torch.from_numpy(image)
    return image

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 create_demo() -> gr.Blocks:
    with gr.Blocks() as demo:
        croper = gr.State()
        with gr.Row():
            with gr.Column():
                edit_prompt = gr.Textbox(lines=1, label="Edit Prompt", value=DEFAULT_EDIT_PROMPT)
                generate_size = gr.Number(label="Generate Size", value=512)
            with gr.Column():
                num_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Num Steps")
                guidance_scale = gr.Slider(minimum=0, maximum=30, value=5, step=0.5, label="Guidance Scale")
            with gr.Column():
                strength = gr.Slider(minimum=0, maximum=3, value=0.2, step=0.1, label="Strength")
                with gr.Accordion("Advanced Options", open=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")
                    seed = gr.Number(label="Seed", value=8)
                    category = gr.Textbox(label="Category", value=DEFAULT_CATEGORY, visible=False)
                    cond_scale1 = gr.Slider(minimum=0, maximum=3, value=0.8, step=0.1, label="Cond_scale1")
                    cond_scale2 = gr.Slider(minimum=0, maximum=3, value=0.3, step=0.1, label="Cond_scale2")
                    lineart_detect = gr.Slider(minimum=0, maximum=1, value=0.375, step=0.01, label="Lineart Detect")
                    canny_detect = gr.Slider(minimum=0, maximum=1, value=0.75, step=0.01, label="Canny Detect")
                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)
            with gr.Column():
                origin_area_image = gr.Image(label="Origin Area Image", type="pil", interactive=False)
                generated_image = gr.Image(label="Generated Image", type="pil", 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, croper],
        ).success(
            fn=image_to_image,
            inputs=[origin_area_image, edit_prompt,seed, num_steps, guidance_scale, strength, generate_size, cond_scale1, cond_scale2, lineart_detect, canny_detect],
            outputs=[generated_image, generated_cost],
        ).success(
            fn=restore_result,
            inputs=[croper, category, generated_image],
            outputs=[restored_image],
        )

    return demo