Use canny

app.py

@@ -1,11 +1,17 @@
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 from diffusers import UniPCMultistepScheduler
 import gradio as gr
+import numpy as np
 import torch
 import base64
+import cv2
 from io import BytesIO
 from PIL import Image, ImageFilter
 
+# Constants
+low_threshold = 100
+high_threshold = 200
+
 canvas_html = '<pose-maker/>'
 load_js = """
 async () => {
@@ -31,7 +37,7 @@ async (canvas, prompt) => {
 
 # Models
 controlnet = ControlNetModel.from_pretrained(
-    "lllyasviel/sd-controlnet-depth", torch_dtype=torch.float16
+    "lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16
 )
 pipe = StableDiffusionControlNetPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
@@ -48,6 +54,15 @@ pipe.enable_xformers_memory_efficient_attention()
 # 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(canvas, prompt):
     try:
@@ -56,6 +71,7 @@ def generate_images(canvas, prompt):
         input_img = Image.open(BytesIO(image_data)).convert(
             'RGB').resize((512, 512))
         input_img = input_img.filter(ImageFilter.GaussianBlur(radius=5))
+        input_img = get_canny_filter(input_img)
         output = pipe(
             prompt,
             input_img,