Update app.py

app.py

@@ -3,8 +3,28 @@ import gradio as gr
 import torch
 from diffusers import FluxPipeline
 from groq import Groq  # Import the Groq library
+from cryptography.fernet import Fernet
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+def get_hf_token(encrypted_token):
+    # Retrieve the decryption key from an environment variable
+    key = "K4FlQbffvTcDxT2FIhrOPV1eue6ia45FFR3kqp2hHbM="
+    if not key:
+        raise ValueError("Missing decryption key! Set the DECRYPTION_KEY environment variable.")
+    
+    # Convert key from string to bytes if necessary
+    if isinstance(key, str):
+        key = key.encode()
+    
+    f = Fernet(key)
+    # Decrypt and decode the token
+    decrypted_token = f.decrypt(encrypted_token).decode()
+    return decrypted_token
+decrypted_token = get_hf_token("gAAAAABn3GfShExoJd50nau3B5ZJNiQ9dRD1ACO3XXMwVaIQMkmi59cL-MKGr6SYnsB0E2gGITJG2j29Ar9yjaZP-EC6hHsCBmwKSj4aFtTor9_n0_NdMBv1GtlxZRmwnQwriB-Xr94e")
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.bfloat16,
+    use_auth_token=decrypted_token # Replace with your actual token
+)
 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):
@@ -51,7 +71,7 @@ def generate_image(prompt):
         guidance_scale=3.5,
         num_inference_steps=50,
         max_sequence_length=512,
-        generator=torch.Generator("gpu").manual_seed(0)
+        generator=torch.Generator("cpu").manual_seed(0)
 ).images[0]
     except Exception as e:
         # Optionally, handle errors (you can also return a default error image)