import streamlit as st
import torch
from transformers import CLIPProcessor, CLIPTextModel  # Using CLIP as an example of a lighter model for image generation

@st.cache_resource
def load_lightweight_model():
    """Load a lightweight image generation model."""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model_name = "CompVis/ldm-text2im-large-256"  # A known lightweight diffusion model
    model = CLIPTextModel.from_pretrained(model_name).to(device)
    processor = CLIPProcessor.from_pretrained(model_name)
    return model, processor

def generate_image(model, processor, prompt, negative_prompt):
    """Generate an image using the lightweight model."""
    try:
        inputs = processor(text=prompt, return_tensors="pt").to(model.device)
        outputs = model.generate(inputs['input_ids'], max_length=50)
        # Here we would simulate or replace this with the actual image output from the lightweight model
        dummy_image = "https://via.placeholder.com/256"  # Placeholder image link for demonstration
        return dummy_image
    except Exception as e:
        st.error(f"Error generating image: {str(e)}")
        return None

# Streamlit app setup
st.title("Image Generator")
st.markdown("### By Taizun", unsafe_allow_html=True)
st.markdown('<small>Effortlessly create stunning images with our lightweight generator.</small>', unsafe_allow_html=True)

prompt = st.text_input("Enter your prompt")
model, processor = load_lightweight_model()

# Style table (Example styles)
styles_dict = {
    "Neon Vibes": {"prompt": "neon lights", "negative_prompt": "blurry, low quality"},
    "Retro Sci-Fi": {"prompt": "retro futuristic scene", "negative_prompt": "modern, dull"},
    "Mystic Forest": {"prompt": "dark forest with mystical lights", "negative_prompt": "bright, plain"},
    "Abstract Art": {"prompt": "abstract shapes and colors", "negative_prompt": "realistic, boring"},
}

# Dropdown for selecting a style
style_name = st.selectbox("Select a Style", options=list(styles_dict.keys()))

# Display the selected style's prompt and negative prompt
if style_name:
    selected_entry = styles_dict[style_name]
    selected_style_prompt = selected_entry["prompt"]
    selected_style_negative_prompt = selected_entry["negative_prompt"]

if st.button("Generate Image"):
    with st.spinner("Generating your image..."):
        result = generate_image(model, processor, prompt + " " + selected_style_prompt, selected_style_negative_prompt)
        if result:
            st.image(result, caption="Generated Image", use_column_width=True)