Update app.py

app.py

@@ -79,13 +79,12 @@ args = Namespace(
     controlnet_revision=None,
     controlnet_from_pt=False,
 )
-weight_dtype = jnp.float32
 
 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(
@@ -93,7 +92,7 @@ pipeline, pipeline_params = FlaxStableDiffusionControlNetPipeline.from_pretraine
     # tokenizer=tokenizer,
     controlnet=controlnet,
     safety_checker=None,
-    dtype=weight_dtype,
+    dtype=jnp.bfloat16,
     revision=args.revision,
     from_pt=args.from_pt,
 )