image2image / app_haircolor_img2img.py
zhiweili
change to controlnet-p2p app
f04f727
raw
history blame
6.23 kB
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 (
StableDiffusionControlNetPipeline,
ControlNetModel,
DDIMScheduler,
DPMSolverMultistepScheduler,
EulerAncestralDiscreteScheduler,
UniPCMultistepScheduler,
)
from controlnet_aux import (
CannyDetector,
LineartDetector,
PidiNetDetector,
HEDdetector,
)
BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DEFAULT_EDIT_PROMPT = "change hair to blue"
DEFAULT_CATEGORY = "hair"
canny_detector = CannyDetector()
lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
lineart_detector = lineart_detector.to(DEVICE)
pidiNet_detector = PidiNetDetector.from_pretrained('lllyasviel/Annotators')
pidiNet_detector = pidiNet_detector.to(DEVICE)
hed_detector = HEDdetector.from_pretrained('lllyasviel/Annotators')
hed_detector = hed_detector.to(DEVICE)
controlnet = [
ControlNetModel.from_pretrained(
"lllyasviel/control_v11e_sd15_ip2p",
torch_dtype=torch.float16,
),
ControlNetModel.from_pretrained(
"lllyasviel/control_v11p_sd15_lineart",
torch_dtype=torch.float16,
),
]
basepipeline = StableDiffusionControlNetPipeline.from_pretrained(
BASE_MODEL,
torch_dtype=torch.float16,
use_safetensors=True,
controlnet=controlnet,
)
basepipeline.scheduler = UniPCMultistepScheduler.from_config(basepipeline.scheduler.config)
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,
image_guidance_scale: float,
generate_size: int,
cond_scale1: float = 1.2,
cond_scale2: float = 1.2,
):
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, 384, generate_size)
cond_image = [input_image, lineart_image]
generator = torch.Generator(device=DEVICE).manual_seed(seed)
generated_image = basepipeline(
generator=generator,
prompt=edit_prompt,
image=cond_image,
height=generate_size,
width=generate_size,
guidance_scale=guidance_scale,
image_guidance_scale=image_guidance_scale,
num_inference_steps=num_steps,
controlnet_conditioning_scale=[cond_scale1, cond_scale2],
).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")
image_guidance_scale = gr.Slider(minimum=0, maximum=30, value=1.5, step=0.1, label="Image Guidance Scale")
with gr.Column():
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=1.2, step=0.1, label="Cond_scale1")
cond_scale2 = gr.Slider(minimum=0, maximum=3, value=1.2, step=0.1, label="Cond_scale2")
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, image_guidance_scale, generate_size, cond_scale1, cond_scale2],
outputs=[generated_image, generated_cost],
).success(
fn=restore_result,
inputs=[croper, category, generated_image],
outputs=[restored_image],
)
return demo