Use input image for ControlNet

app.py

@@ -4,39 +4,37 @@ import numpy as np
 import torch
 
 from diffusers import DiffusionPipeline, StableDiffusionControlNetPipeline
-from diffusers.utils import load_image
 from PIL import Image
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-image = load_image("https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png")
-image = np.array(image)
-
 low_threshold = 100
 high_threshold = 200
 
-image = cv2.Canny(image, low_threshold, high_threshold)
-image = image[:, :, None]
-image = np.concatenate([image, image, image], axis=2)
-canny_image = Image.fromarray(image)
-
 def generate(
-    prompt, negative_prompt, num_inference_steps, width, height, guidance_scale, seed
+    prompt, negative_prompt, num_inference_steps, width, height, guidance_scale, seed, input_image
 ):
+    if input_image is None:
+        pipeline = DiffusionPipeline.from_pretrained("Lykon/DreamShaper")
+        pipeline = pipeline.to(device)
+        generator = torch.Generator(device=device).manual_seed(seed)
+
+        return pipeline(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        ).images[0]
+
+    image = cv2.Canny(input_image, low_threshold, high_threshold)
+    image = image[:, :, None]
+    image = np.concatenate([image, image, image], axis=2)
+    canny_image = Image.fromarray(image)
+
     return canny_image
-    # pipeline = DiffusionPipeline.from_pretrained("Lykon/DreamShaper")
-    # pipeline = pipeline.to(device)
-    # generator = torch.Generator(device=device).manual_seed(seed)
-
-    # return pipeline(
-    #     prompt=prompt,
-    #     negative_prompt=negative_prompt,
-    #     num_inference_steps=num_inference_steps,
-    #     width=width,
-    #     height=height,
-    #     guidance_scale=guidance_scale,
-    #     generator=generator,
-    # ).images[0]
 
 iface = gr.Interface(
     fn=generate,
@@ -54,6 +52,7 @@ iface = gr.Interface(
             step=1,
             randomize=True,
         ),
+        gr.Image(label="Input Image", source='upload', type="numpy")
     ],
     outputs="image",
 )