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

# Title and Description
st.title("Denim Virtual Try-On")
st.write("Upload your leg images and the front/back images of denim pants to see how they fit.")

# File Upload Section
st.sidebar.header("Upload Images")
subject_image = st.sidebar.file_uploader("Upload Subject Image (Legs)", type=["jpg", "jpeg", "png"])
pant_front_image = st.sidebar.file_uploader("Upload Denim Pant Front Image", type=["jpg", "jpeg", "png"])
pant_back_image = st.sidebar.file_uploader("Upload Denim Pant Back Image", type=["jpg", "jpeg", "png"])

# Function to process images using Hugging Face API
def virtual_tryon(subject_img, front_img, back_img):
    url = "https://api-inference.huggingface.co/models/ramim36/Kolors-Virtual-Try-On"
    headers = {"Authorization": f"Bearer {os.getenv('HF_API_TOKEN')}"}  # Set your Hugging Face API token in env variable
    
    # Prepare data
    files = {
        "subject_image": subject_img,
        "front_image": front_img,
        "back_image": back_img
    }
    
    response = requests.post(url, headers=headers, files=files)
    if response.status_code == 200:
        return Image.open(BytesIO(response.content))
    else:
        st.error(f"Error: {response.status_code} - {response.text}")
        return None

# Display Input Images
if subject_image:
    st.image(subject_image, caption="Uploaded Subject Image", use_column_width=True)
if pant_front_image:
    st.image(pant_front_image, caption="Uploaded Front Denim Image", use_column_width=True)
if pant_back_image:
    st.image(pant_back_image, caption="Uploaded Back Denim Image", use_column_width=True)

# Generate Try-On Output
if st.button("Generate Virtual Try-On"):
    if subject_image and pant_front_image and pant_back_image:
        with st.spinner("Processing... Please wait."):
            try:
                composite_image = virtual_tryon(subject_image, pant_front_image, pant_back_image)
                if composite_image:
                    st.image(composite_image, caption="Virtual Try-On Result", use_column_width=True)
            except Exception as e:
                st.error(f"An error occurred: {e}")
    else:
        st.error("Please upload all required images.")