import streamlit as st
from diffusers import DiffusionPipeline
import torch

# Set page config
st.set_page_config(
    page_title="Portrait Generator",
    page_icon="🖼️",
    layout="centered"
)

# App title and description
st.title("AI Portrait Generator")
st.markdown("Generate beautiful portraits using the AWPortraitCN2 model")

# Model parameters
with st.sidebar:
    st.header("Generation Settings")
    steps = st.slider("Inference Steps", min_value=20, max_value=100, value=40)
    guidance_scale = st.slider("Guidance Scale", min_value=1.0, max_value=15.0, value=7.5, step=0.5)
    negative_prompt = st.text_area(
        "Negative Prompt",
        value="lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, watermark, signature, out of frame"
    )
    seed = st.number_input("Random Seed (leave at -1 for random)", min_value=-1, value=-1)
    
# Main prompt input
prompt = st.text_area(
    "Describe the portrait you want to generate",
    value="Masterpiece portrait of a beautiful young woman with flowing hair, detailed face, photorealistic, 8k, professional photography"
)

# Generate button
if st.button("Generate Portrait", type="primary"):
    with st.spinner("Loading model and generating portrait..."):
        try:
            # Set up the model pipeline
            pipeline = DiffusionPipeline.from_pretrained(
                "Shakker-Labs/AWPortraitCN2",
                torch_dtype=torch.float16,
                use_safetensors=True
            )
            
            # Move to GPU if available
            device = "cuda" if torch.cuda.is_available() else "cpu"
            pipeline = pipeline.to(device)
            
            # Set seed if specified
            generator = None
            if seed != -1:
                generator = torch.Generator(device).manual_seed(seed)
            
            # Generate the image
            image = pipeline(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=steps,
                guidance_scale=guidance_scale,
                generator=generator
            ).images[0]
            
            # Display the generated image
            st.image(image, caption="Generated Portrait", use_column_width=True)
            
            # Option to download
            # Convert the PIL image to bytes
            import io
            from PIL import Image
            buf = io.BytesIO()
            image.save(buf, format="PNG")
            byte_im = buf.getvalue()
            
            st.download_button(
                label="Download Portrait",
                data=byte_im,
                file_name="generated_portrait.png",
                mime="image/png"
            )
            
        except Exception as e:
            st.error(f"An error occurred: {e}")
            st.info("Make sure you have enough GPU memory and the required dependencies installed.")

# Add requirements info at the bottom
st.markdown("---")
st.markdown("""
### Requirements
To run this app, you need:
- diffusers
- transformers
- accelerate
- torch
- streamlit

Install with: `pip install diffusers transformers accelerate torch streamlit`
""")