import streamlit as st
import torch
import numpy as np
from diffusers import DiffusionPipeline
from transformers import pipeline

# Load text generation pipeline
@st.cache_resource
def load_text_pipeline():
    return pipeline('text-generation', model='daspartho/prompt-extend')

text_pipe = load_text_pipeline()

def extend_prompt(prompt):
    return text_pipe(prompt + ',', num_return_sequences=1, max_new_tokens=50)[0]["generated_text"]

@st.cache_resource
def load_pipeline():
    pipe = DiffusionPipeline.from_pretrained(
        "stabilityai/sdxl-turbo",
        torch_dtype=torch.float32,
        variant=None,
        use_safetensors=True
    )
    pipe.to("cpu")
    return pipe

def generate_image(prompt, use_details):
    pipe = load_pipeline()
    generator = torch.manual_seed(np.random.randint(0, 2**32))
    extended_prompt = extend_prompt(prompt) if use_details else prompt
    image = pipe(prompt=extended_prompt, generator=generator, num_inference_steps=15, guidance_scale=7.5).images[0]
    return image, extended_prompt

# Add the custom CSS file
st.markdown('<link rel="stylesheet" href="styles.css">', unsafe_allow_html=True)

st.markdown("<div class='header'>✨ Generate Your Custom GitHub Profile Picture! ✨</div>", unsafe_allow_html=True)
st.markdown("<div class='subheader'>Create an anime-style GitHub profile picture that reflects your personality and passion for coding. 🚀👨‍💻</div>", unsafe_allow_html=True)

# Input widgets
input_text = st.text_area("Describe your GitHub profile picture:", "Create an anime-style GitHub profile picture for a boy")

details_checkbox = st.checkbox("Generate Details?", value=True)

if st.button("Generate Image"):
    with st.spinner('Generating image...'):
        image, extended_prompt = generate_image(input_text, details_checkbox)
        st.image(image, caption="Generated Image")
        st.text(f"Extended Prompt: {extended_prompt}")
    st.balloons()