from diffusers import StableDiffusionControlNetPipeline, ControlNetModel from diffusers import UniPCMultistepScheduler import cv2 import gradio as gr import numpy as np import torch from PIL import Image # Constants low_threshold = 100 high_threshold = 200 # Models controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16) pipe = StableDiffusionControlNetPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16 ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.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.enable_model_cpu_offload() # Generator seed, generator = torch.manual_seed(0) def get_canny_filter(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 generate_images(image, prompt): canny_image = get_canny_filter(image) output = pipe( prompt, canny_image, generator=generator, num_images_per_prompt=3 ) all_outputs = [] all_outputs.append(canny_image) for image in output.images: all_outputs.append(image) return all_outputs gr.Interface( generate_images, inputs=[ gr.Image(type="pil"), gr.Textbox( label="Enter your prompt", max_lines=1, placeholder="Sandra Oh, best quality, extremely detailed", ), ], outputs=gr.Gallery().style(grid=[2], height="auto"), title="Generate controlled outputs with ControlNet and Stable Diffusion. ", description="This Space uses Canny edge maps as the additional conditioning.", examples=[["input_image_vermeer.png", "Sandra Oh, best quality, extremely detailed"]], allow_flagging=False, ).launch(enable_queue=True)