import streamlit as st
from diffusers import DiffusionPipeline
import torch
import os

@st.cache_resource
def load_pipeline():
    # Get the token from the environment variable
    token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if not token:
        st.error("Hugging Face token not found. Please check your Hugging Face Spaces secrets.")
        st.stop()

    pipeline = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", use_auth_token=token)
    pipeline.load_lora_weights("gorkemyurt/lora-train")
    return pipeline

st.title("FLUX.1 Diffusion Model with LoRA")

pipeline = load_pipeline()

prompt = st.text_input("Enter your prompt:", "A beautiful landscape with mountains and a lake")
num_inference_steps = st.slider("Number of inference steps:", min_value=1, max_value=100, value=50)
guidance_scale = st.slider("Guidance scale:", min_value=1.0, max_value=20.0, value=7.5, step=0.1)

if st.button("Generate Image"):
    with st.spinner("Generating image..."):
        image = pipeline(
            prompt=prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale
        ).images[0]
    
    st.image(image, caption="Generated Image", use_column_width=True)