import cv2 import gradio as gr import numpy as np import torch from diffusers import ( ControlNetModel, StableDiffusionControlNetPipeline, UniPCMultistepScheduler, ) from PIL import Image stable_model_list = [ "runwayml/stable-diffusion-v1-5", "stabilityai/stable-diffusion-2-1", ] controlnet_canny_model_list = [ "lllyasviel/sd-controlnet-canny", "thibaud/controlnet-sd21-canny-diffusers", ] stable_prompt_list = ["a photo of a man.", "a photo of a girl."] stable_negative_prompt_list = ["bad, ugly", "deformed"] data_list = [ "data/test.png", ] def controlnet_canny( dict_image: str, controlnet_model_path: str, ): image = dict_image["image"].convert("RGB").resize((512, 512)) image = np.array(image) image = cv2.Canny(image, 100, 200) image = image[:, :, None] image = np.concatenate([image, image, image], axis=2) image = Image.fromarray(image) controlnet = ControlNetModel.from_pretrained( controlnet_model_path, torch_dtype=torch.float16 ) return controlnet, image def stable_diffusion_controlnet_canny( image_path: str, stable_model_path: str, controlnet_model_path: str, prompt: str, negative_prompt: str, guidance_scale: int, num_inference_step: int, ): controlnet, image = controlnet_canny( image_path=image_path, controlnet_model_path=controlnet_model_path ) pipe = StableDiffusionControlNetPipeline.from_pretrained( pretrained_model_name_or_path=stable_model_path, controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16, ) pipe.to("cuda") pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) pipe.enable_xformers_memory_efficient_attention() output = pipe( prompt=prompt, image=image, negative_prompt=negative_prompt, num_inference_steps=num_inference_step, guidance_scale=guidance_scale, ).images return output[0] def stable_diffusion_controlnet_canny_app(): with gr.Blocks(): with gr.Row(): with gr.Column(): controlnet_canny_image_file = gr.Image( type="filepath", label="Image" ) controlnet_canny_stable_model_id = gr.Dropdown( choices=stable_model_list, value=stable_model_list[0], label="Stable Model Id", ) controlnet_canny_model_id = gr.Dropdown( choices=controlnet_canny_model_list, value=controlnet_canny_model_list[0], label="Controlnet Model Id", ) controlnet_canny_prompt = gr.Textbox( lines=1, value=stable_prompt_list[0], label="Prompt" ) controlnet_canny_negative_prompt = gr.Textbox( lines=1, value=stable_negative_prompt_list[0], label="Negative Prompt", ) with gr.Accordion("Advanced Options", open=False): controlnet_canny_guidance_scale = gr.Slider( minimum=0.1, maximum=15, step=0.1, value=7.5, label="Guidance Scale", ) controlnet_canny_num_inference_step = gr.Slider( minimum=1, maximum=100, step=1, value=50, label="Num Inference Step", ) controlnet_canny_predict = gr.Button(value="Generator") with gr.Column(): output_image = gr.Image(label="Output") gr.Examples( fn=stable_diffusion_controlnet_canny, examples=[ [ data_list[0], stable_model_list[0], controlnet_canny_model_list[0], stable_prompt_list[0], stable_negative_prompt_list[0], 7.5, 50, ] ], inputs=[ controlnet_canny_image_file, controlnet_canny_stable_model_id, controlnet_canny_model_id, controlnet_canny_prompt, controlnet_canny_negative_prompt, controlnet_canny_guidance_scale, controlnet_canny_num_inference_step, ], outputs=[output_image], cache_examples=False, label="Controlnet Canny Example", ) controlnet_canny_predict.click( fn=stable_diffusion_controlnet_canny, inputs=[ controlnet_canny_image_file, controlnet_canny_stable_model_id, controlnet_canny_model_id, controlnet_canny_prompt, controlnet_canny_negative_prompt, controlnet_canny_guidance_scale, controlnet_canny_num_inference_step, ], outputs=[output_image], )