updated model, added canny filter too

app.py

@@ -21,7 +21,7 @@ low_threshold = 100
 high_threshold = 200
 
 base_model_path = "runwayml/stable-diffusion-v1-5"
-controlnet_path = "lmattingly/controlnet-uncanny-simpsons"
+controlnet_path = "lmattingly/controlnet-uncanny-simpsons-v2-0"
 #controlnet_path = "JFoz/dog-cat-pose"
 
 # Models
@@ -29,9 +29,29 @@ controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
     controlnet_path, dtype=jnp.bfloat16
 )
 pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, revision="flax", dtype=jnp.bfloat16
+    base_model_path, controlnet=controlnet, revision="flax", dtype=jnp.bfloat16
 )
 
+
+def canny_filter(image):
+    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    blurred_image = cv2.GaussianBlur(gray_image, (5, 5), 0)
+    edges_image = cv2.Canny(blurred_image, 50, 150)
+    canny_image = Image.fromarray(edges_image)
+    return canny_image
+
+def canny_filter2(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)
+    return canny_image 
+
+
+
 def resize_image(im, max_size):
     im_np = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
 
@@ -45,19 +65,20 @@ def resize_image(im, max_size):
 
     return resized_im
 
+
 def create_key(seed=0):
     return jax.random.PRNGKey(seed)
 
 def infer(prompts, image):
     params["controlnet"] = controlnet_params
-    
     im = image
-    image = resize_image(im, 500)
+    image = canny_filter2(im)
+    #image = canny_filter(im)
+    #image = Image.fromarray(im)
+    
     num_samples = 1 #jax.device_count()
     rng = create_key(0)
     rng = jax.random.split(rng, jax.device_count())
-    #im = image
-    #image = Image.fromarray(im)
     
     prompt_ids = pipe.prepare_text_inputs([prompts] * num_samples)
     processed_image = pipe.prepare_image_inputs([image] * num_samples)