import streamlit as st
from PIL import Image
import torch

# Optimized lightweight image generation (example)
@st.cache_resource
def load_lightweight_model():
    # Assuming there's a lightweight model for CPU
    # Placeholder: Replace this with actual lightweight model loading
    return "lightweight_model"

model = load_lightweight_model()

# Streamlit app
st.title("Text-to-Image Generation (Fast)")

# User input for prompt
user_prompt = st.text_input("Enter your image prompt", value="Astronaut in a jungle")

# Button to generate the image
if st.button("Generate Image"):
    if user_prompt:
        with st.spinner("Generating image..."):
            # Simulated fast image generation (substitute with lightweight model logic)
            # Example: Generate a simple placeholder image
            image = Image.new("RGB", (256, 256), color="blue")  # Placeholder image
            
            # Display the generated image
            st.image(image, caption="Generated Image", use_column_width=True)
    else:
        st.error("Please enter a valid prompt.")