import streamlit as st
from diffusers import DiffusionPipeline
from PIL import Image
import torch

# Load the FLUX model
@st.cache_resource
def load_pipeline():
    # Using the 'black-forest-labs/FLUX.1-schnell' model
    pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
    pipe.enable_model_cpu_offload()  # Offload to CPU to save memory
    return pipe

pipe = load_pipeline()

# Streamlit app
st.title("Text-to-Image Generation with FLUX.1-schnell")

# User input for prompt
user_prompt = st.text_input("Enter your image prompt", value="A cat holding a sign that says hello world")

# Button to generate the image
if st.button("Generate Image"):
    if user_prompt:
        with st.spinner("Generating image..."):
            # Generate the image using the FLUX model
            image = pipe(
                user_prompt,
                guidance_scale=0.0,  # No guidance
                num_inference_steps=4,  # Number of steps for faster generation
            ).images[0]
            
            # Save and display the image
            image.save("generated_image.png")
            st.image(image, caption="Generated Image (FLUX.1-schnell)", use_column_width=True)
    else:
        st.error("Please enter a valid prompt.")