hatimanees commited on
Commit
e2d8327
·
verified ·
1 Parent(s): 0033c30

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +22 -14
app.py CHANGED
@@ -1,24 +1,25 @@
1
  import streamlit as st
2
  from PIL import Image
3
  from diffusers import DiffusionPipeline
 
4
  import os
5
 
6
- # Hugging Face Authentication
7
  @st.cache_resource
8
  def authenticate_and_load_model(hf_token):
9
  """
10
- Authenticate with Hugging Face and load the model.
11
  """
12
- os.environ["HF_HOME"] = "./hf_cache" # Set cache directory for Hugging Face models
13
- os.environ["HF_AUTH_TOKEN"] = hf_token # Set authentication token for Hugging Face
14
-
15
  try:
 
 
 
16
  # Load the model
17
  pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", use_auth_token=hf_token)
18
  pipe.load_lora_weights("tryonlabs/FLUX.1-dev-LoRA-Lehenga-Generator", use_auth_token=hf_token)
19
  return pipe
20
  except Exception as e:
21
- st.error(f"Error loading model: {e}")
22
  return None
23
 
24
  # Streamlit App
@@ -27,19 +28,26 @@ st.write("Enter a description to generate an image of a lehenga dress.")
27
 
28
  # Hugging Face Token Input
29
  hf_token = st.text_input("Enter your Hugging Face Token:", type="password")
 
 
30
  if hf_token:
31
- pipe = authenticate_and_load_model(hf_token)
32
- else:
33
- st.warning("Please enter your Hugging Face token to proceed.")
 
34
 
35
  # Input prompt
36
- prompt = st.text_area("Enter your prompt:",
37
- "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.")
 
 
38
 
39
  # Generate button
40
  if st.button("Generate Image"):
41
- if not hf_token or not pipe:
42
- st.error("Model not loaded. Please ensure you've provided a valid Hugging Face token.")
 
 
43
  elif prompt.strip():
44
  with st.spinner("Generating image..."):
45
  try:
@@ -48,7 +56,7 @@ if st.button("Generate Image"):
48
  # Display the image
49
  st.image(result, caption="Generated Lehenga Image", use_column_width=True)
50
  except Exception as e:
51
- st.error(f"An error occurred: {e}")
52
  else:
53
  st.warning("Please enter a valid prompt.")
54
 
 
1
  import streamlit as st
2
  from PIL import Image
3
  from diffusers import DiffusionPipeline
4
+ from huggingface_hub import login
5
  import os
6
 
7
+ # Hugging Face Login Function
8
  @st.cache_resource
9
  def authenticate_and_load_model(hf_token):
10
  """
11
+ Log in to Hugging Face and load the model with LoRA weights.
12
  """
 
 
 
13
  try:
14
+ # Log in to Hugging Face
15
+ login(token=hf_token)
16
+
17
  # Load the model
18
  pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", use_auth_token=hf_token)
19
  pipe.load_lora_weights("tryonlabs/FLUX.1-dev-LoRA-Lehenga-Generator", use_auth_token=hf_token)
20
  return pipe
21
  except Exception as e:
22
+ st.error(f"Error during login or model loading: {e}")
23
  return None
24
 
25
  # Streamlit App
 
28
 
29
  # Hugging Face Token Input
30
  hf_token = st.text_input("Enter your Hugging Face Token:", type="password")
31
+ pipe = None
32
+
33
  if hf_token:
34
+ if "pipe" not in st.session_state:
35
+ with st.spinner("Authenticating and loading the model..."):
36
+ st.session_state.pipe = authenticate_and_load_model(hf_token)
37
+ pipe = st.session_state.pipe
38
 
39
  # Input prompt
40
+ prompt = st.text_area(
41
+ "Enter your prompt:",
42
+ "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."
43
+ )
44
 
45
  # Generate button
46
  if st.button("Generate Image"):
47
+ if not hf_token:
48
+ st.error("Please enter your Hugging Face token.")
49
+ elif not pipe:
50
+ st.error("Model not loaded. Please check your Hugging Face token.")
51
  elif prompt.strip():
52
  with st.spinner("Generating image..."):
53
  try:
 
56
  # Display the image
57
  st.image(result, caption="Generated Lehenga Image", use_column_width=True)
58
  except Exception as e:
59
+ st.error(f"An error occurred during image generation: {e}")
60
  else:
61
  st.warning("Please enter a valid prompt.")
62