Update app.py

app.py

@@ -17,7 +17,7 @@ using our implementation of the RAFT model. We will also see how to convert the
 predicted flows to RGB images for visualization.
 """
 
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers import DiffusionPipeline, ControlNetModel
 from diffusers import UniPCMultistepScheduler
 
 import cv2
@@ -48,7 +48,7 @@ high_threshold = 200
 
 # Models
 controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
-pipe = StableDiffusionControlNetPipeline.from_pretrained(
+pipe = DiffusionPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None, torch_dtype=torch.float16
 )
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
@@ -73,11 +73,12 @@ def get_canny_filter(image):
     return canny_image
 
 
-def generate_images(prompt, canny_image):
+def generate_images(prompt, canny_image, image):
     
     output = pipe(
-        prompt,
-        canny_image,
+        #prompt,
+        controlnet_conditioning_image=canny_image,
+        image = image,
         generator=generator,
         num_images_per_prompt=1,
         num_inference_steps=20,
@@ -128,7 +129,7 @@ def infer():
     
     pil2diff_img = Image.open("./basket1.jpg")
     canny_image = get_canny_filter(pil2diff_img)
-    diffused_img = generate_images(prompt, canny_image)
+    diffused_img = generate_images(prompt, canny_image, pil2diff_img)
     print(f"DIFFUSED IMG: {diffused_img[1]}")
     
     diffused_img[1].save("diffused_input1.jpg")
@@ -275,7 +276,7 @@ def infer():
 
     pil2diff_blend = Image.open("blended2.jpg")
     canny_image = get_canny_filter(pil2diff_blend)
-    diffused_blend = generate_images(prompt, canny_image)
+    diffused_blend = generate_images(prompt, canny_image, pil2diff_blend)
     print(f"DIFFUSED IMG: {diffused_blend[1]}")
     
     diffused_blend[1].save("diffused_blended_2.jpg")