Update app.py

app.py

@@ -1,7 +1,7 @@
 from PIL import Image
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import FluxPipeline
 from groq import Groq  # Import the Groq library
 from cryptography.fernet import Fernet
 from huggingface_hub import login
@@ -25,9 +25,12 @@ def get_hf_token(encrypted_token):
 groq_client = Groq(api_key="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")
 decrypted_token = get_hf_token("gAAAAABn3GfShExoJd50nau3B5ZJNiQ9dRD1ACO3XXMwVaIQMkmi59cL-MKGr6SYnsB0E2gGITJG2j29Ar9yjaZP-EC6hHsCBmwKSj4aFtTor9_n0_NdMBv1GtlxZRmwnQwriB-Xr94e")
 login(token=decrypted_token)
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell")
+
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", 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
 
+
 def enhance_prompt(user_prompt):
     """Enhances the given prompt using Groq and returns the refined prompt."""
     try: