import streamlit as st
from PIL import Image
from io import BytesIO
from transformers import pipeline

# Title and App Description
st.title("👗 Virtual Dress Try-On")
st.write("""
Upload a **Human Body Image** and a **Garment Image** to generate a Virtual Try-On.  
Images will be compressed to meet Hugging Face size constraints.
""")

# Function to load and compress images
def compress_image(image_file, max_size_kb=512):
    """
    Compresses the input image to meet size constraints while preserving the aspect ratio.
    Args:
        image_file: Uploaded file from Streamlit
        max_size_kb: Maximum file size in kilobytes
    Returns:
        Compressed PIL Image object
    """
    img = Image.open(image_file).convert("RGB")
    quality = 95  # Initial compression quality
    img_format = "JPEG"

    # Compress image iteratively until it meets size constraints
    while True:
        img_bytes = BytesIO()
        img.save(img_bytes, format=img_format, quality=quality)
        size_kb = len(img_bytes.getvalue()) / 1024  # Size in KB

        if size_kb <= max_size_kb or quality <= 10:
            break
        quality -= 5  # Reduce quality to compress further

    compressed_img = Image.open(img_bytes)
    return compressed_img

# Load Model (Hugging Face Pipeline)
@st.cache_resource
def load_model():
    model_pipeline = pipeline("image-to-image", model="ares1123/virtual-dress-try-on")
    return model_pipeline

model = load_model()

# Sidebar for Image Upload
st.sidebar.header("Upload Images")
uploaded_person = st.sidebar.file_uploader("Upload Human Body Image", type=["jpg", "jpeg", "png"])
uploaded_clothing = st.sidebar.file_uploader("Upload Garment Image", type=["jpg", "jpeg", "png"])

# Process and Display Images
if uploaded_person and uploaded_clothing:
    # Compress uploaded images
    st.sidebar.info("Compressing images to meet size constraints...")
    person_image = compress_image(uploaded_person)
    garment_image = compress_image(uploaded_clothing)

    # Display compressed images
    col1, col2 = st.columns(2)
    col1.subheader("Compressed Human Body Image")
    col1.image(person_image, use_column_width=True)

    col2.subheader("Compressed Garment Image")
    col2.image(garment_image, use_column_width=True)

    # Process button
    if st.button("👗 Generate Virtual Try-On"):
        with st.spinner("Processing images... Please wait ⏳"):
            try:
                # Prepare inputs for the model
                inputs = {"image": person_image, "clothing": garment_image}
                
                # Generate output using Hugging Face model
                output_image = model(inputs)
                
                # Display output image
                st.subheader("✨ Virtual Try-On Result")
                st.image(output_image, use_column_width=True, caption="Composite Virtual Try-On Image")
            except Exception as e:
                st.error(f"An error occurred during processing: {e}")
else:
    st.warning("Please upload both the Human Body Image and Garment Image.")

# Footer
st.markdown("---")
st.write("Developed with ❤️ using Streamlit and Hugging Face.")