Update app.py

app.py

@@ -36,9 +36,6 @@ def create_key(seed=0):
     return jax.random.PRNGKey(seed)
 
 
-rng = create_key(0)
-
-
 
 
 def canny_filter(image):
@@ -48,19 +45,21 @@ def canny_filter(image):
     edges_image = cv2.Canny(blurred_image, 50, 150)
     return edges_image
 
+# load control net and stable diffusion v1-5
+controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
+    "jax-diffusers-event/canny-coyo1m", dtype=jnp.float32
+)
+pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, from_pt=True, dtype=jnp.float32
+)
+
 def infer(prompts, negative_prompts, image):
     
     
-    # load control net and stable diffusion v1-5
-    controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
-        "jax-diffusers-event/canny-coyo1m", dtype=jnp.float32
-    )
-    pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
-        "runwayml/stable-diffusion-v1-5", controlnet=controlnet, from_pt=True, dtype=jnp.float32
-    )
     params["controlnet"] = controlnet_params
     
     num_samples = jax.device_count()
+    rng = create_key(0)
     rng = jax.random.split(rng, jax.device_count())
     canny_image = canny_filter(image)