app.py

import streamlit as st
import os
from vton import virtual_try_on
import requests
from PIL import Image
from io import BytesIO

# Set page config for a wider layout and custom title
st.set_page_config(
    page_title="Virtual Try-On App",
    page_icon="👔",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .stApp {
        max-width: 1200px;
        margin: 0 auto;
    }
    .upload-box {
        border: 2px dashed #4c4c4c;
        border-radius: 10px;
        padding: 20px;
        text-align: center;
        margin: 10px 0;
    }
    .output-box {
        border: 1px solid #e0e0e0;
        border-radius: 10px;
        padding: 20px;
        margin: 20px 0;
        background-color: #f8f9fa;
    }
</style>
""", unsafe_allow_html=True)

# Title and description
st.title("✨ Virtual Try-On Application")
st.markdown("""Upload a garment image and a person's photo to see how the garment would look on them!""")

# Create two columns for input images
col1, col2 = st.columns(2)

# Garment image upload
with col1:
    st.subheader("Garment Image")
    garment_source = st.radio(
        "Choose garment image source:",
        ["Upload", "URL"],
        key="garment_source"
    )
    
    if garment_source == "Upload":
        garment_file = st.file_uploader("Upload garment image", type=["jpg", "jpeg", "png", "webp"], key="garment_upload")
        if garment_file:
            st.image(garment_file, caption="Uploaded Garment", use_container_width=True)
            # Save uploaded file to temporary URL
            garment_img = Image.open(garment_file)
            garment_buffer = BytesIO()
            garment_img.save(garment_buffer, format="PNG")
            garment_url = "temp_garment.png"
            with open(garment_url, "wb") as f:
                f.write(garment_buffer.getvalue())
    else:
        garment_url = st.text_input("Enter garment image URL", key="garment_url")
        if garment_url:
            try:
                st.image(garment_url, caption="Garment from URL", use_container_width=True)
            except:
                st.error("Unable to load image from URL")

# Person image upload
with col2:
    st.subheader("Person Image")
    person_source = st.radio(
        "Choose person image source:",
        ["Upload", "URL"],
        key="person_source"
    )
    
    if person_source == "Upload":
        person_file = st.file_uploader("Upload person image", type=["jpg", "jpeg", "png", "webp"], key="person_upload")
        if person_file:
            st.image(person_file, caption="Uploaded Person", use_container_width=True)
            # Save uploaded file to temporary URL
            person_img = Image.open(person_file)
            person_buffer = BytesIO()
            person_img.save(person_buffer, format="PNG")
            person_url = "temp_person.png"
            with open(person_url, "wb") as f:
                f.write(person_buffer.getvalue())
    else:
        person_url = st.text_input("Enter person image URL", key="person_url")
        if person_url:
            try:
                st.image(person_url, caption="Person from URL", use_container_width=True)
            except:
                st.error("Unable to load image from URL")

# Garment description input
st.subheader("Garment Description")
garment_desc = st.text_area(
    "Describe the garment (e.g., color, style, type)",
    height=100,
    placeholder="Example: A cute pink top with floral pattern"
)

# Check for API token
if "REPLICATE_API_TOKEN" not in os.environ:
    st.warning(
        "⚠️ REPLICATE_API_TOKEN is not set. Please set it using: "
        "export REPLICATE_API_TOKEN='your_token_here'"
    )

# Process button
if st.button("Generate Try-On", type="primary"):
    if not (garment_url and person_url and garment_desc):
        st.error("Please provide all required inputs!")
    else:
        try:
            with st.spinner("🔄 Processing virtual try-on..."):
                output_path = virtual_try_on(garment_url, person_url, garment_desc)
                
                # Display result in a styled box
                st.markdown("<div class='output-box'>", unsafe_allow_html=True)
                st.subheader("🎉 Try-On Result")
                st.image(output_path, caption="Virtual Try-On Result", use_container_width=True)
                st.markdown("</div>", unsafe_allow_html=True)
                
                # Cleanup temporary files
                if os.path.exists("temp_garment.png"):
                    os.remove("temp_garment.png")
                if os.path.exists("temp_person.png"):
                    os.remove("temp_person.png")
                    
        except Exception as e:
            st.error(f"Error occurred: {str(e)}")

# Footer
st.markdown("---")
st.markdown("""
<div style='text-align: center; color: #666;'>
    <p>Powered by IDM-VTON model via Replicate API</p>
</div>
""", unsafe_allow_html=True)