import streamlit as st
from PIL import Image
import os
import tempfile
from genai import gen_vton
from werkzeug.utils import secure_filename

# Define a function to save the uploaded file to a temporary directory
def save_uploaded_file(uploaded_file):
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix="." + uploaded_file.name.split('.')[1]) as tmp:
            tmp.write(uploaded_file.getbuffer())
            return tmp.name
    except Exception as e:
        return None

# Streamlit UI components
st.title('Image Processing with gen_vton')

# File uploader for the user and dress images
user_image_up = st.file_uploader("Upload User Image", type=['jpg', 'png', 'jpeg'])
dress_image_up = st.file_uploader("Upload Dress Image", type=['jpg', 'png', 'jpeg'])

# Process the images once both are uploaded and a button is pressed
if user_image_up and dress_image_up and st.button('Process Images'):
    user_image_path = save_uploaded_file(user_image_up)
    dress_image_path = save_uploaded_file(dress_image_up)

    # Assuming gen_vton processes the images and saves them somewhere
    processed_image_paths = gen_vton(user_image_path, dress_image_path)  # Adjust based on actual return value

    # Displaying processed images (assuming gen_vton returns paths to processed images)
    if processed_image_paths:
        for image_path in processed_image_paths:
            st.image(image_path, use_column_width=True)

    # Clean up temporary files
    os.remove(user_image_path)
    os.remove(dress_image_path)