Update app.py

app.py

@@ -1,27 +1,46 @@
 import streamlit as st
 from PIL import Image
 from diffusers import DiffusionPipeline
+import os
 
-# Load the model and LoRA weights
+# Hugging Face Authentication
 @st.cache_resource
-def load_model():
-    pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev")
-    pipe.load_lora_weights("tryonlabs/FLUX.1-dev-LoRA-Lehenga-Generator")
-    return pipe
+def authenticate_and_load_model(hf_token):
+    """
+    Authenticate with Hugging Face and load the model.
+    """
+    os.environ["HF_HOME"] = "./hf_cache"  # Set cache directory for Hugging Face models
+    os.environ["HF_AUTH_TOKEN"] = hf_token  # Set authentication token for Hugging Face
+    
+    try:
+        # Load the model
+        pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", use_auth_token=hf_token)
+        pipe.load_lora_weights("tryonlabs/FLUX.1-dev-LoRA-Lehenga-Generator", use_auth_token=hf_token)
+        return pipe
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        return None
 
-pipe = load_model()
-
-# App interface
+# Streamlit App
 st.title("Lehenga Dress Image Generator")
 st.write("Enter a description to generate an image of a lehenga dress.")
 
+# Hugging Face Token Input
+hf_token = st.text_input("Enter your Hugging Face Token:", type="password")
+if hf_token:
+    pipe = authenticate_and_load_model(hf_token)
+else:
+    st.warning("Please enter your Hugging Face token to proceed.")
+
 # Input prompt
 prompt = st.text_area("Enter your prompt:", 
                       "A flat-lay image of a lehenga with a traditional style and a fitted waistline is elegantly crafted from stretchy silk material, ensuring a comfortable and flattering fit. The long hemline adds a touch of grace and sophistication to the ensemble. Adorned in a solid blue color, it features a sleeveless design that complements its sweetheart neckline. The solid pattern and the luxurious silk fabric together create a timeless and chic look that is perfect for special occasions.")
 
 # Generate button
 if st.button("Generate Image"):
-    if prompt.strip():
+    if not hf_token or not pipe:
+        st.error("Model not loaded. Please ensure you've provided a valid Hugging Face token.")
+    elif prompt.strip():
         with st.spinner("Generating image..."):
             try:
                 # Generate the image