import streamlit as st
from diffusers import StableDiffusionPipeline
import torch

@st.cache_resource
def load_lightweight_pipeline():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipeline = StableDiffusionPipeline.from_pretrained(
        "CompVis/stable-diffusion-v1-4", 
        torch_dtype=torch.float16 if device == "cuda" else torch.float32
    ).to(device)
    return pipeline

def generate_image(pipeline, prompt, negative_prompt):
    try:
        image = pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=25,
            guidance_scale=7.5
        ).images[0]
        return image
    except Exception as e:
        st.error(f"Error generating image: {str(e)}")
        return None

st.title("Image Generator")
st.markdown('<small>By Taizun</small>', unsafe_allow_html=True)

prompt = st.text_input("Enter your prompt")
pipeline = load_lightweight_pipeline()

styles_dict = {
    "Neon Vibes": {"prompt": "neon lights in a futuristic city", "negative_prompt": "blurred, dull"},
    "Retro Sci-Fi": {"prompt": "retro futuristic landscape", "negative_prompt": "realistic, modern"},
    "Mystic Forest": {"prompt": "dark mystical forest with glowing lights", "negative_prompt": "sunny, bright"},
    "Abstract Art": {"prompt": "abstract colorful shapes", "negative_prompt": "realistic, detailed"}
}

style_name = st.selectbox("Select a Style", options=list(styles_dict.keys()))

if style_name:
    selected_style = styles_dict[style_name]
    style_prompt = selected_style["prompt"]
    negative_prompt = selected_style["negative_prompt"]

if st.button("Generate Image"):
    with st.spinner("Generating your image..."):
        image = generate_image(pipeline, prompt + " " + style_prompt, negative_prompt)
        if image:
            st.image(image, caption="Generated Image", use_column_width=True)