Update app.py

app.py

@@ -51,6 +51,23 @@ def nms(x, t, s):
     z[y > t] = 255
     return z
 
+def HWC3(x):
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y
 
 DESCRIPTION = '''# Scribble SDXL 🖋️🌄
 sketch to image with SDXL, using [@xinsir](https://huggingface.co/xinsir) [scribble sdxl controlnet](https://huggingface.co/xinsir/controlnet-scribble-sdxl-1.0)
@@ -131,7 +148,10 @@ controlnet = ControlNetModel.from_pretrained(
     "xinsir/controlnet-scribble-sdxl-1.0",
     torch_dtype=torch.float16
 )
-
+controlnet_canny = ControlNetModel.from_pretrained(
+    "xinsir/controlnet-canny-sdxl-1.0",
+    torch_dtype=torch.float16
+)
 # when test with other base model, you need to change the vae also.
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
@@ -144,6 +164,16 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
 )
 pipe.to(device)
 # Load model.
+pipe_canny = StableDiffusionXLControlNetPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    controlnet=controlnet_canny,
+    vae=vae,
+    safety_checker=None,
+    torch_dtype=torch.float16,
+    scheduler=eulera_scheduler,
+)
+
+pipe_canny.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 processor = HEDdetector.from_pretrained('lllyasviel/Annotators')
@@ -180,6 +210,7 @@ def run(
     controlnet_conditioning_scale: float = 1.0,
     seed: int = 0,
     use_hed: bool = False,
+    use_canny: bool = False,
     progress=gr.Progress(track_tqdm=True),
 ) -> PIL.Image.Image:
     width, height  = image['composite'].size
@@ -187,7 +218,13 @@ def run(
     new_width, new_height = int(width * ratio), int(height * ratio)
     image = image['composite'].resize((new_width, new_height))
     
-    if not use_hed:
+    if use_canny:
+        controlnet_img = np.array(image)
+        controlnet_img = cv2.Canny(controlnet_img, 100, 200)
+        controlnet_img = HWC3(controlnet_img)
+        image = Image.fromarray(controlnet_img)
+    
+    elif not use_hed:
           controlnet_img = image
     else:
         controlnet_img = processor(image, scribble=False)
@@ -206,7 +243,8 @@ def run(
     prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
 
     generator = torch.Generator(device=device).manual_seed(seed)
-    out = pipe(
+    if use_canny:
+        out = pipe_canny(
         prompt=prompt,
         negative_prompt=negative_prompt,
         image=image,
@@ -217,6 +255,17 @@ def run(
         width=new_width,
         height=new_height,
     ).images[0]
+    else:
+        out = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=image,
+        num_inference_steps=num_steps,
+        generator=generator,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        guidance_scale=guidance_scale,
+        width=new_width,
+        height=new_height,).images[0]
 
     return (controlnet_img, out)
 
@@ -236,6 +285,7 @@ with gr.Blocks(css="style.css", js=js_func) as demo:
                 prompt = gr.Textbox(label="Prompt")
                 style = gr.Dropdown(label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
                 use_hed = gr.Checkbox(label="use HED detector", value=False, info="check this box if you upload an image and want to turn it to a sketch")
+                use_canny = gr.Checkbox(label="use Canny", value=False, info="check this to use ControlNet canny instead of scribble")
                 run_button = gr.Button("Run")
             with gr.Accordion("Advanced options", open=False):
                 negative_prompt = gr.Textbox(
@@ -287,6 +337,7 @@ with gr.Blocks(css="style.css", js=js_func) as demo:
         controlnet_conditioning_scale,
         seed,
         use_hed,
+        use_canny
     ]
     outputs = [image_slider]
     run_button.click(