diffusers/controlnet-canny

1from diffusers import StableDiffusionControlNetPipeline, ControlNetModel2from diffusers import UniPCMultistepScheduler3import cv24import gradio as gr5import numpy as np6import torch7from PIL import Image8 9# Constants10low_threshold = 10011high_threshold = 20012 13# Models14controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)15pipe = StableDiffusionControlNetPipeline.from_pretrained(16    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float1617)18pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)19 20# This command loads the individual model components on GPU on-demand. So, we don't21# need to explicitly call pipe.to("cuda").22pipe.enable_model_cpu_offload()23 24pipe.enable_xformers_memory_efficient_attention()25 26# Generator seed,27generator = torch.manual_seed(0)28 29def get_canny_filter(image):30    if not isinstance(image, np.ndarray):31        image = np.array(image) 32        33    image = cv2.Canny(image, low_threshold, high_threshold)34    image = image[:, :, None]35    image = np.concatenate([image, image, image], axis=2)36    canny_image = Image.fromarray(image)37    return canny_image38 39 40def generate_images(image, prompt):41    canny_image = get_canny_filter(image)42    output = pipe(43        prompt,44        canny_image,45        generator=generator,46        num_images_per_prompt=3,47        num_inference_steps=20,48    )49    all_outputs = []50    all_outputs.append(canny_image)51    for image in output.images:52        all_outputs.append(image)53    return all_outputs54 55 56gr.Interface(57    generate_images,58    inputs=[59        gr.Image(type="pil"),60        gr.Textbox(61            label="Enter your prompt",62            max_lines=1,63            placeholder="Sandra Oh, best quality, extremely detailed",64        ),65    ],66    outputs=gr.Gallery().style(grid=[2], height="auto"),67    title="Generate controlled outputs with ControlNet and Stable Diffusion. ",68    description="This Space uses Canny edge maps as the additional conditioning.",69    examples=[["input_image_vermeer.png", "Sandra Oh, best quality, extremely detailed"]],70    allow_flagging=False,71).launch(enable_queue=True)72 73 74