Update app.py

app.py

@@ -28,14 +28,14 @@ def create_key(seed=0):
 
 # load control net and stable diffusion v1-5
 controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
-    "learner/jax-diffuser-event", from_flax=True, dtype=jnp.float32
+    "learner/jax-diffuser-event", from_flax=True, dtype=jnp.bfloat16
 )
 
 pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
     controlnet=controlnet,
     from_pt=True,
-    dtype=jnp.float32,
+    dtype=jnp.bfloat16,
     #safety_checker=None,
 )