Update app.py

app.py

@@ -7,15 +7,11 @@ from diffusers import FluxControlPipeline, FluxTransformer2DModel
 ####################################
 #   Load the model(s) on CPU       #
 ####################################
-edit_transformer = FluxTransformer2DModel.from_pretrained(
-    "sayakpaul/FLUX.1-dev-edit-v0", 
-    torch_dtype=torch.float32
-)
+path = "sayakpaul/FLUX.1-dev-edit-v0" 
+edit_transformer = FluxTransformer2DModel.from_pretrained(path, torch_dtype=torch.bfloat16)
 pipeline = FluxControlPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev",
-    transformer=edit_transformer,
-    torch_dtype=torch.float32
-)
+    "black-forest-labs/FLUX.1-dev", transformer=edit_transformer, torch_dtype=torch.bfloat16
+).to("cuda")
 
 #####################################
 #  The function for our Gradio app  #