Update app.py

app.py

@@ -1,10 +1,11 @@
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import FluxPipeline
 from PIL import Image
 
 # Load the diffusion model
-pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev")
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
 
 # Set the model to the appropriate device
 device = "cuda" if torch.cuda.is_available() else "cpu"