Update app.py

app.py

@@ -84,7 +84,7 @@ controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
     args.controlnet_model_name_or_path,
     revision=args.controlnet_revision,
     from_pt=args.controlnet_from_pt,
-    dtype=jnp.float32,
+    dtype=jnp.bfloat16,
 )
 
 pipeline, pipeline_params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
@@ -92,7 +92,7 @@ pipeline, pipeline_params = FlaxStableDiffusionControlNetPipeline.from_pretraine
     # tokenizer=tokenizer,
     controlnet=controlnet,
     safety_checker=None,
-    dtype=jnp.float32,
+    dtype=jnp.bfloat16,
     revision=args.revision,
     from_pt=args.from_pt,
 )