sayakpaul HF staff commited on
Commit
ae4a684
1 Parent(s): cf6cfad

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +64 -0
app.py CHANGED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
2
+ from diffusers import UniPCMultistepScheduler
3
+ from diffusers.utils import load_image
4
+ import gradio as gr
5
+ import torch
6
+
7
+ # Constants
8
+ low_threshold = 100
9
+ high_threshold = 200
10
+
11
+ # Models
12
+ controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
13
+ pipe = StableDiffusionControlNetPipeline.from_pretrained(
14
+ "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
15
+ )
16
+ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
17
+
18
+ # This command loads the individual model components on GPU on-demand. So, we don't
19
+ # need to explicitly call pipe.to("cuda").
20
+ pipe.enable_model_cpu_offload()
21
+
22
+ # Generator seed,
23
+ generator = torch.manual_seed(0)
24
+
25
+ def get_canny_filter(image):
26
+ if not isinstance(image, np.ndarray):
27
+ image = np.array(image)
28
+
29
+ image = cv2.Canny(image, low_threshold, high_threshold)
30
+ image = image[:, :, None]
31
+ image = np.concatenate([image, image, image], axis=2)
32
+ canny_image = Image.fromarray(image)
33
+ return canny_image
34
+
35
+
36
+ def generate_images(image, prompt):
37
+ canny_image = get_canny_filter(image)
38
+ output = pipe(
39
+ prompt,
40
+ canny_image,
41
+ generator=generator,
42
+ num_images_per_prompt=3
43
+ )
44
+ return output.images
45
+
46
+
47
+ gr.Interface(
48
+ generate_images,
49
+ inputs=[
50
+ gr.Image(type="pil"),
51
+ gr.Textbox(
52
+ label="Enter your prompt",
53
+ max_lines=1,
54
+ placeholder="Sandra Oh, best quality, extremely detailed",
55
+ ),
56
+ ],
57
+ outputs=gr.Gallery().style(grid=[2], height="auto"),
58
+ title="Generate controlled outputs with ControlNet and Stable Diffusion. ",
59
+ description="This Space uses Canny edge maps as the additional conditioning.",
60
+ examples=[["input_image_vermeer.png", "Sandra Oh, best quality, extremely detailed"]],
61
+ allow_flagging=False,
62
+ ).launch(enable_queue=True)
63
+
64
+