Update app.py

app.py

@@ -1,11 +1,17 @@
 import os
 import random
 import numpy as np
-import spaces  # Import spaces before torch
+import spaces
 import torch
 from diffusers import FluxPipeline
 import gradio as gr
 
+# Access the Hugging Face token from environment variables
+hf_token = os.getenv("HF_TOKEN")
+
+if hf_token is None:
+    raise ValueError("Hugging Face token is not set. Please set the HF_TOKEN environment variable.")
+
 # Check if GPU is available
 if torch.cuda.is_available():
     device = "cuda"
@@ -14,14 +20,15 @@ else:
     device = "cpu"
     print("Using CPU")
 
-# login hf token
-HF_TOKEN = os.getenv("HF_TOKEN")
-
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
 
 # Initialize the pipeline and download the model
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-dev",
+    torch_dtype=torch.bfloat16,
+    use_auth_token=hf_token  # Use the token from the environment variable
+)
 pipe.to(device)
 
 # Enable memory optimizations