Spaces:
Sleeping
Sleeping
File size: 2,844 Bytes
93f033c e2d8327 552867f 93f033c 640c05a e2d8327 93f033c 552867f 640c05a 552867f e2d8327 640c05a 552867f e2d8327 552867f 93f033c 552867f 93f033c 552867f e2d8327 552867f e2d8327 552867f 93f033c e2d8327 93f033c e2d8327 552867f 93f033c e2d8327 93f033c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
import streamlit as st
from PIL import Image
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, download the model, and save it locally for reuse.
"""
try:
# Log in to Hugging Face
login(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}")
return None
# 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")
pipe = None
if hf_token:
if "pipe" not in st.session_state:
with st.spinner("Authenticating and loading the model..."):
st.session_state.pipe = authenticate_and_load_model(hf_token)
pipe = st.session_state.pipe
# 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 not hf_token:
st.error("Please enter your Hugging Face token.")
elif not pipe:
st.error("Model not loaded. Please check your Hugging Face token.")
elif prompt.strip():
with st.spinner("Generating image..."):
try:
# Generate the image
result = pipe(prompt).images[0]
# Display the image
st.image(result, caption="Generated Lehenga Image", use_column_width=True)
except Exception as e:
st.error(f"An error occurred during image generation: {e}")
else:
st.warning("Please enter a valid prompt.")
st.write("This app uses AI to generate images of lehenga dresses based on your input description.")
|