import streamlit as st
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
import numpy as np

# Set page config
st.set_page_config(
    page_title="AI Image Generator",
    page_icon="🎨",
    layout="centered"
)

# Cache the model loading to avoid reloading on every interaction
@st.cache_resource
def load_model():
    """Load and cache the Stable Diffusion model"""
    model_id = "runwayml/stable-diffusion-v1-5"
    
    # Check if CUDA is available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Load the pipeline
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        use_safetensors=True
    )
    pipe = pipe.to(device)
    
    # Enable memory efficient attention if using CUDA
    if device == "cuda":
        pipe.enable_attention_slicing()
        pipe.enable_memory_efficient_attention()
    
    return pipe

def generate_image(prompt, negative_prompt="", num_inference_steps=20, guidance_scale=7.5, width=512, height=512):
    """Generate image from text prompt"""
    try:
        pipe = load_model()
        
        # Generate image
        with torch.no_grad():
            image = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                width=width,
                height=height
            ).images[0]
        
        return image
    except Exception as e:
        st.error(f"Error generating image: {str(e)}")
        return None

def main():
    # Header
    st.title("🎨 AI Image Generator")
    st.markdown("Generate beautiful images from text descriptions using Stable Diffusion!")
    
    # Sidebar for advanced settings
    with st.sidebar:
        st.header("⚙️ Settings")
        
        # Image dimensions
        col1, col2 = st.columns(2)
        with col1:
            width = st.selectbox("Width", [512, 768, 1024], index=0)
        with col2:
            height = st.selectbox("Height", [512, 768, 1024], index=0)
        
        # Generation parameters
        num_inference_steps = st.slider("Inference Steps", 10, 50, 20, 
                                       help="More steps = better quality but slower")
        guidance_scale = st.slider("Guidance Scale", 1.0, 20.0, 7.5, 0.5,
                                  help="Higher values = more adherence to prompt")
        
        # Info
        st.markdown("---")
        st.markdown("### 💡 Tips")
        st.markdown("- Be specific in your descriptions")
        st.markdown("- Use artistic styles (e.g., 'oil painting', 'digital art')")
        st.markdown("- Add quality modifiers (e.g., 'highly detailed', '4k')")
        st.markdown("- Use negative prompts to avoid unwanted elements")
    
    # Main content area
    col1, col2 = st.columns([2, 1])
    
    with col1:
        # Text input for prompt
        prompt = st.text_area(
            "✍️ Describe the image you want to generate:",
            placeholder="A beautiful sunset over mountains, oil painting style, highly detailed",
            height=100
        )
        
        # Negative prompt (optional)
        negative_prompt = st.text_area(
            "❌ Negative prompt (optional - things to avoid):",
            placeholder="blurry, low quality, distorted",
            height=60
        )
        
        # Generate button
        generate_btn = st.button("🚀 Generate Image", type="primary", use_container_width=True)
    
    with col2:
        # Example prompts
        st.markdown("### 🎯 Example Prompts")
        examples = [
            "A majestic lion in a savanna at sunset",
            "Cyberpunk cityscape at night, neon lights",
            "Van Gogh style painting of a coffee shop",
            "Cute robot playing with cats in a garden",
            "Abstract art with vibrant colors and geometric shapes"
        ]
        
        for i, example in enumerate(examples):
            if st.button(f"Use Example {i+1}", key=f"example_{i}"):
                st.session_state.example_prompt = example
        
        # Apply example if selected
        if hasattr(st.session_state, 'example_prompt'):
            prompt = st.session_state.example_prompt
            del st.session_state.example_prompt
            st.rerun()
    
    # Generate and display image
    if generate_btn and prompt:
        with st.spinner("🎨 Creating your masterpiece... This may take a few moments!"):
            # Show progress
            progress_bar = st.progress(0)
            for i in range(100):
                progress_bar.progress(i + 1)
                if i == 99:
                    break
            
            # Generate image
            image = generate_image(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                width=width,
                height=height
            )
            
            progress_bar.empty()
        
        if image:
            # Display the generated image
            st.success("✅ Image generated successfully!")
            st.image(image, caption=f"Generated from: '{prompt}'", use_column_width=True)
            
            # Download button
            img_buffer = io.BytesIO()
            image.save(img_buffer, format='PNG')
            st.download_button(
                label="📥 Download Image",
                data=img_buffer.getvalue(),
                file_name=f"generated_image_{hash(prompt) % 10000}.png",
                mime="image/png",
                use_container_width=True
            )
            
            # Show generation parameters
            with st.expander("📊 Generation Details"):
                st.json({
                    "prompt": prompt,
                    "negative_prompt": negative_prompt,
                    "dimensions": f"{width}x{height}",
                    "inference_steps": num_inference_steps,
                    "guidance_scale": guidance_scale
                })
    
    elif generate_btn and not prompt:
        st.warning("⚠️ Please enter a prompt to generate an image!")
    
    # Footer
    st.markdown("---")
    st.markdown(
        "Built with ❤️ using [Streamlit](https://streamlit.io) and "
        "[Stable Diffusion](https://huggingface.co/runwayml/stable-diffusion-v1-5)"
    )

if __name__ == "__main__":
    # Add missing import
    import io
    main()