import streamlit as st
import torch
from diffusers import DiffusionPipeline

# Model paths and settings
color_book_lora_path = "artificialguybr/ColoringBookRedmond-V2"
color_book_trigger = ", ColoringBookAF, Coloring Book"

@st.cache_resource
def load_pipeline(lora):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    st.write(f"Using device: {device}")  # Displaying the selected device
    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", 
                                             torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                                             use_safetensors=True,
                                             variant="fp16" if device == "cuda" else None)
    if lora != "None":
        st.write(f"Loading Lora model: {color_book_lora_path}")  # Debugging Lora loading
        pipe.load_lora_weights(color_book_lora_path)
    return pipe

def image_generation(pipe, prompt, negative_prompt):
    try:
        image = pipe(
            prompt=prompt,
            negative_prompt="blurred, ugly, watermark, low resolution" + negative_prompt,
            num_inference_steps=20,
            guidance_scale=9.0
        ).images[0]
        return image
    except Exception as e:
        st.error(f"Error generating image: {str(e)}")
        return None

# Data for different styles
table = [
    {
        "name": "sai-neonpunk",
        "prompt": "neonpunk style . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
        "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured"
    },
    {
        "name": "futuristic-retro cyberpunk",
        "prompt": "retro cyberpunk. 80's inspired, synthwave, neon, vibrant, detailed, retro futurism",
        "negative_prompt": "modern, desaturated, black and white, realism, low contrast"
    },
    {
        "name": "Dark Fantasy",
        "prompt": "Dark Fantasy Art, dark, moody, dark fantasy style",
        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, bright, sunny"
    },
    {
        "name": "Double Exposure",
        "prompt": "Double Exposure Style, double image ghost effect, image combination, double exposure style",
        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
    },
    {
        "name": "None",
        "prompt": "8K ",
        "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured"
    }
]

# Convert list to dict for easier lookup
styles_dict = {entry["name"]: entry for entry in table}

# Streamlit app title and styling
st.title("Coloring Book Generator")
st.markdown("<h2 style='text-align: center; color: #888;'>By Taizun</h2>", unsafe_allow_html=True)

# Add a header image
st.image("https://your-image-link.com", use_container_width=True)  # Replace with your image URL or local path

# Prompt input from user
prompt = st.text_input("Enter your Prompt", value="A cute Lion")

# Dropdown for selecting Lora
select_lora = st.selectbox("Select your Lora", options=["Coloring Book", "None"])

# Dropdown for selecting 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"]

# Button for image generation
if st.button("Generate Image"):
    with st.spinner("Generating your awesome image..."):
        pipeline = load_pipeline(select_lora)
        if select_lora == "None":
            image = image_generation(pipeline, prompt + selected_style_prompt, selected_style_negative_prompt)
        else:
            image = image_generation(pipeline, prompt + selected_style_prompt + color_book_trigger, selected_style_negative_prompt)
        if image:
            st.image(image)

# Add a footer with credit by Taizun (smaller text)
st.markdown("""
    <style>
        .footer {
            text-align: center;
            font-size: 0.8rem;
            color: #888;
            margin-top: 30px;
        }
    </style>
    <div class="footer">
        <p>Created by Taizun | Powered by Streamlit</p>
    </div>
    """, unsafe_allow_html=True)