Update app.py

app.py

@@ -51,9 +51,26 @@ def nms(x, t, s):
     z[y > t] = 255
     return z
 
-
-DESCRIPTION = '''# ⚡️Flash⚡️ Scribble SDXL 🖋️🌄
-super fast sketch to image with Flash SDXL, using [@xinsir](https://huggingface.co/xinsir) [scribble sdxl controlnet](https://huggingface.co/xinsir/controlnet-scribble-sdxl-1.0) and [sdxl flash](https://huggingface.co/sd-community/sdxl-flash)
+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)
 '''
 
 if not torch.cuda.is_available():
@@ -124,11 +141,21 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# eulera_scheduler = EulerAncestralDiscreteScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
+
+
 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)
 
@@ -140,7 +167,21 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
     # scheduler=eulera_scheduler,
 )
 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+
 pipe.to(device)
+
+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.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe_canny.scheduler.config)
+
+pipe_canny.to(device)
 # Load model.
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -178,6 +219,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
@@ -185,7 +227,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)
@@ -204,7 +252,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,
@@ -215,6 +264,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)
 
@@ -234,6 +294,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(
@@ -243,9 +304,9 @@ with gr.Blocks(css="style.css", js=js_func) as demo:
                 num_steps = gr.Slider(
                     label="Number of steps",
                     minimum=1,
-                    maximum=20,
+                    maximum=50,
                     step=1,
-                    value=10,
+                    value=1,
                 )
                 guidance_scale = gr.Slider(
                     label="Guidance scale",
@@ -285,6 +346,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(