import streamlit as st
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from huggingface_hub import login
from tqdm.auto import tqdm
import os
from PIL import Image
import io

class ProgressCallback:
    def __init__(self, progress_bar):
        self.progress_bar = progress_bar
        
    def __call__(self, step: int, timestep: int, latents: torch.FloatTensor):
        self.progress_bar.progress(step / self.progress_bar.total)

def setup_stable_diffusion(hf_token):
    """
    Setup Stable Diffusion with proper authentication and GPU optimization
    
    Args:
        hf_token (str): HuggingFace API token
    """
    # Check GPU availability
    if not torch.cuda.is_available():
        st.warning("WARNING: No GPU detected. Processing will be slow on CPU.")
        return None

    # Display GPU info
    st.info(f"Using GPU: {torch.cuda.get_device_name()}")
    st.info(f"GPU Memory Available: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")

    # Login to Hugging Face
    login(token=hf_token)
    
    # Initialize the pipeline
    model_id = "runwayml/stable-diffusion-v1-5"
    
    with st.spinner("Loading Stable Diffusion model..."):
        pipe = StableDiffusionPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            use_safetensors=True,
            token=hf_token
        )
        
        # Use DPM++ 2M scheduler
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        
        # Move to GPU and optimize
        pipe = pipe.to("cuda")
        pipe.enable_attention_slicing()
        pipe.enable_vae_slicing()
        
        # Try to enable xformers
        try:
            pipe.enable_xformers_memory_efficient_attention()
            st.success("Successfully enabled xformers memory efficient attention")
        except Exception:
            st.warning("Note: xformers not available. Install with: pip install xformers")
        
        pipe.safety_checker = None
    
    return pipe

def generate_image(pipe, prompt, num_inference_steps=30):
    """
    Generate an image using the initialized pipeline
    
    Args:
        pipe: Initialized StableDiffusionPipeline
        prompt (str): Text prompt for image generation
        num_inference_steps (int): Number of denoising steps
    """
    if pipe is None:
        st.error("Error: Pipeline not properly initialized")
        return None

    # Clear CUDA cache before generation
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    # Create a progress bar
    progress_bar = st.progress(0)
    status_text = st.empty()
    
    # Update the callback to use Streamlit's progress bar
    pipe.callback = ProgressCallback(progress_bar)
    
    status_text.text(f"Generating image for prompt: '{prompt}'")
    
    with torch.autocast("cuda"):
        image = pipe(
            prompt=prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=7.5,
        ).images[0]
    
    progress_bar.empty()
    status_text.empty()
    
    return image

def main():
    st.title("🎨 Stable Diffusion Image Generator")
    
    # Initialize session state
    if 'pipe' not in st.session_state:
        st.session_state.pipe = None
    if 'model_initialized' not in st.session_state:
        st.session_state.model_initialized = False
    
    # Sidebar for configuration
    st.sidebar.header("Configuration")
    
    # HuggingFace token input
    hf_token = st.sidebar.text_input("Enter HuggingFace Token:", type="password")
    
    # Auto-initialize when token is entered
    if hf_token and not st.session_state.model_initialized:
        st.session_state.pipe = setup_stable_diffusion(hf_token)
        if st.session_state.pipe is not None:
            st.session_state.model_initialized = True
            st.sidebar.success("Model initialized successfully!")
    
    # Main interface
    prompt = st.text_input("Enter your prompt:", "A stunning landscape with mountains and a lake at sunset")
    num_steps = st.slider("Number of inference steps:", min_value=20, max_value=100, value=30)
    
    if st.button("Generate Image"):
        if not st.session_state.model_initialized:
            st.error("Please enter your HuggingFace token in the sidebar to initialize the model!")
        else:
            try:
                image = generate_image(st.session_state.pipe, prompt, num_steps)
                if image is not None:
                    st.image(image, caption=prompt, use_column_width=True)
                    
                    # Add download button
                    buf = io.BytesIO()
                    image.save(buf, format='PNG')
                    st.download_button(
                        label="Download Image",
                        data=buf.getvalue(),
                        file_name=f"generated_image_{prompt[:30]}.png",
                        mime="image/png"
                    )
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

if __name__ == "__main__":
    main()