import streamlit as st
from PIL import Image
import os
# import torch

style_dir = "imgs/style"
content_dir = "imgs/content"
results_dir = "imgs/results"
if not os.path.exists(style_dir):
    os.mkdir(style_dir)
if not os.path.exists(content_dir):
    os.mkdir(content_dir)
if not os.path.exists(results_dir):
    os.mkdir(results_dir)

st.set_page_config(page_title="DiffStyle Demo", page_icon=":art:", layout="wide")


style_reference_image = None
content_reference_image = None

# Get style reference image
style_reference_image = st.sidebar.file_uploader("Upload style reference image", type=["jpg", "jpeg", "png"], key='style_reference')
if style_reference_image:
    style_reference_image = Image.open(style_reference_image)
    st.sidebar.image(style_reference_image, width=200)

# Get content reference image
content_reference_image = st.sidebar.file_uploader("Upload content reference image", type=["jpg", "jpeg", "png"], key='content_reference')
if content_reference_image:
    content_reference_image = Image.open(content_reference_image)
    st.sidebar.image(content_reference_image, width=200)

# Synthesize new image
# if style_reference_image and content_reference_image:
#     # Call your AI algorithm here to synthesize a new image
#     synthesized_image = None # placeholder for your synthesized image
    
#     st.image(synthesized_image, caption="Synthesized image", use_column_width=True)

st.write("""
This is just a demo. Your AI algorithm code should be added to synthesize the new image.
""")