Update app.py

app.py

@@ -4,19 +4,30 @@ from diffusers import DiffusionPipeline
 from huggingface_hub import login
 import os
 
+# Set a custom directory to save the model
+MODEL_DIR = "./saved_models"
+
 # Hugging Face Login Function
 @st.cache_resource
 def authenticate_and_load_model(hf_token):
     """
-    Log in to Hugging Face and load the model with LoRA weights.
+    Log in to Hugging Face, download the model, and save it locally for reuse.
     """
     try:
         # Log in to Hugging Face
         login(token=hf_token)
 
-        # 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)
+        # Load the model and LoRA weights, saving them to the custom directory
+        pipe = DiffusionPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-dev", 
+            cache_dir=MODEL_DIR,
+            use_auth_token=hf_token
+        )
+        pipe.load_lora_weights(
+            "tryonlabs/FLUX.1-dev-LoRA-Lehenga-Generator", 
+            cache_dir=MODEL_DIR,
+            use_auth_token=hf_token
+        )
         return pipe
     except Exception as e:
         st.error(f"Error during login or model loading: {e}")