import cv2
import einops
import gradio as gr
import numpy as np
import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers import UniPCMultistepScheduler
from PIL import Image
from controlnet_aux import OpenposeDetector
# Constants
low_threshold = 100
high_threshold = 200
# Models
controlnet_canny = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
pipe_canny = StableDiffusionControlNetPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", controlnet=controlnet_canny, safety_checker=None, torch_dtype=torch.float16
)
pipe_canny.scheduler = UniPCMultistepScheduler.from_config(pipe_canny.scheduler.config)
# This command loads the individual model components on GPU on-demand. So, we don't
# need to explicitly call pipe.to("cuda").
pipe_canny.enable_model_cpu_offload()
pipe_canny.enable_xformers_memory_efficient_attention()
# Generator seed,
generator = torch.manual_seed(0)
pose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
controlnet_pose = ControlNetModel.from_pretrained(
"lllyasviel/sd-controlnet-openpose", torch_dtype=torch.float16
)
pipe_pose = StableDiffusionControlNetPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", controlnet=controlnet_pose, safety_checker=None, torch_dtype=torch.float16
)
pipe_pose.scheduler = UniPCMultistepScheduler.from_config(pipe_pose.scheduler.config)
# This command loads the individual model components on GPU on-demand. So, we don't
# need to explicitly call pipe.to("cuda").
pipe_pose.enable_model_cpu_offload()
# xformers
pipe_pose.enable_xformers_memory_efficient_attention()
def get_canny_filter(image):
print(image)
if not isinstance(image, np.ndarray):
image = np.array(image)
image = cv2.Canny(image, low_threshold, high_threshold)
image = image[:
, :, None]
image = np.concatenate([image, image, image], axis=2)
canny_image = Image.fromarray(image)
return canny_image
def get_pose(image):
return pose_model(image)
def process(input_image, prompt, input_control):
# TODO: Add other control tasks
if input_control == "Scribble":
return process_canny(input_image, prompt)
elif input_control == "Pose":
return process_pose(input_image, prompt)
return process_canny(input_image, prompt)
def process_canny(input_image, prompt):
canny_image = get_canny_filter(input_image)
output = pipe_canny(
prompt,
canny_image,
generator=generator,
num_images_per_prompt=1,
num_inference_steps=20,
)
return [canny_image,output.images[0]]
def process_pose(input_image, prompt):
pose_image = get_pose(input_image)
output = pipe_pose(
prompt,
pose_image,
generator=generator,
num_images_per_prompt=1,
num_inference_steps=20,
)
return [pose_image,output.images[0]]
block = gr.Blocks().queue()
control_task_list = [
"Canny Edge Map",
"Scribble",
"Pose"
]
with block:
gr.Markdown("## Adding Conditional Control to Text-to-Image Diffusion Models")
gr.HTML('''
This is an unofficial demo for ControlNet, which is a neural network structure to control diffusion models by adding extra conditions such as canny edge detection. The demo is based on the Github implementation.
''')
gr.HTML("You can duplicate this Space to run it privately without a queue and load additional checkpoints. : 
")
with gr.Row():
with gr.Column():
input_image = gr.Image(source='upload', type="numpy")
input_control = gr.Dropdown(control_task_list, value="Scribble", label="Control Task")
prompt = gr.Textbox(label="Prompt")
run_button = gr.Button(label="Run")
with gr.Column():
result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
ips = [input_image, prompt, input_control]
run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
examples_list = [
[
"bird.png",
"bird",
"Canny Edge Map"
],
# [
# "turtle.png",
# "turtle",
# "Scribble",
# "best quality, extremely detailed",
# 'longbody, lowres, bad anatomy, bad hands, missing fingers, pubic hair,extra digit, fewer digits, cropped, worst quality, low quality',
# 1,
# 512,
# 20,
# 9.0,
# 123490213,
# 0.0,
# 100,
# 200
# ],
[
"pose1.png",
"Chef in the Kitchen",
"Pose",
# "best quality, extremely detailed",
# 'longbody, lowres, bad anatomy, bad hands, missing fingers, pubic hair,extra digit, fewer digits, cropped, worst quality, low quality',
# 1,
# 512,
# 20,
# 9.0,
# 123490213,
# 0.0,
# 100,
# 200
]
]
examples = gr.Examples(examples=examples_list,inputs = [input_image, prompt, input_control], outputs = [result_gallery], cache_examples = True, fn = process)
gr.Markdown("")
block.launch(debug = True)