import streamlit as st
import torch
from diffusers import DiffusionPipeline

# Load the Stable Diffusion model on CPU
@st.cache_resource
def load_model():
    try:
        model_id = "CompVis/stable-diffusion-v1-4"  # Smaller model
        pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
        pipe = pipe.to("cpu")
        return pipe
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None

pipe = load_model()

# Streamlit app
st.title("LogoCraft - AI Logo Generator")
st.write("Generate logos using Stable Diffusion and Hugging Face Spaces.")

# Input fields
logo_name = st.text_input("Logo Name or Description", placeholder="e.g., My Awesome Brand")
color_palette = st.selectbox("Color Palette", ["Vibrant", "Pastel", "Monochrome", "Black & White"])
logo_style = st.selectbox("Logo Style", ["Cartoon", "Mascot", "Minimal", "3D"])
effects = st.selectbox("Effects", ["None", "Gradient", "Shadow", "Glow"])
design_idea = st.text_area("Design Idea (Optional)", placeholder="e.g., I like geometric shapes and modern fonts")

# Generate button
if st.button("Generate Logo"):
    if logo_name:
        # Create a detailed prompt
        prompt = (
            f"A logo for {logo_name}, "
            f"using a {color_palette} color palette, "
            f"in a {logo_style} style, "
            f"with {effects} effects, "
            f"inspired by: {design_idea if design_idea else 'modern design trends'}"
        )

        # Generate the logo
        with st.spinner("Generating logo..."):
            try:
                image = pipe(prompt).images[0]
                # Display the logo
                st.image(image, caption="Generated Logo", use_column_width=True)

                # Download button
                st.download_button(
                    label="Download Logo",
                    data=image.tobytes(),
                    file_name="logo.png",
                    mime="image/png"
                )
            except Exception as e:
                st.error(f"Error generating logo: {e}")
    else:
        st.error("Please enter a logo name or description.")