import streamlit as st
from streamlit_drawable_canvas import st_canvas
from PIL import Image
from typing import Union
import random
import numpy as np
import os
import time

st.set_page_config(layout="wide")


def create_edit_existing_image_tab():
    st.write("# Edit existing image")
    

    cols = st.columns(2)
    with cols[0]:
        image_source = st.file_uploader("Upload source image", type=["png", "jpg", "jpeg", "webp"], key="upload_source_edit_existing_image")
        st.text_input("Source object", key="text_input_source_edit_existing_image")
        st.image('content/dog.png')
    with cols[1]:
        image_target = st.file_uploader("Upload target image", type=["png", "jpg", "jpeg", "webp"], key="upload_target_edit_existing_image")
        st.text_input("Target object", key="text_input_target_edit_existing_image")
        st.image('content/cat-sofa.png')

    st.text_input("Prompt", key="text_input_prompt_edit_existing_image")
    st.text_input("Negative prompt", key="text_input_negative_prompt_edit_existing_image")
    st.button("Generate", key="button_generate_edit_existing_image")
    
    st.write("## Result")
    st.image('content/after_editing.png')
    

def create_edit_generated_image_tab():
    st.write("# Edit generated image")
    
    cols = st.columns(2)
    with cols[0]:
        image_source = st.file_uploader("Upload source image", type=["png", "jpg", "jpeg", "webp"], key="upload_source_edit_generated_image")
        st.text_input("Target object", key="text_input_source_edit_generated_image")
        st.text_input("Prompt", key="text_input_prompt_edit_generated_image")
        st.text_input("Negative prompt", key="text_input_negative_prompt_edit_generated_image")
        if image_source:
            st.button("Generate", key="button_generate_edit_generated_image")
    with cols[1]:
        st.image('content/dog.png')

    
    st.write("## Result")
    cols_result = st.columns(2)
    with cols_result[0]:
        st.write("### Generated image before editing")
        st.image('content/before_editing_generated.png')
    with cols_result[1]:
        st.write("### Generated image after editing")
        st.image('content/after_editing_generated.png')
    
def create_zero_shot_generation_tab():
    st.write("# Zero-shot generation")
    

def create_zero_shot_stylization_tab():
    st.write("# Zero-shot stylization")
    
def create_home_tab():
    st.write("# Home of BLIP-Diffusion")
    st.write("Welcome to the demo application of BLIP-Diffusion")
    
    st.write("Project page is [here](https://dxli94.github.io/BLIP-Diffusion-website/.)")
    st.write("Github page is [here](https://github.com/salesforce/LAVIS/tree/main/projects/blip-diffusion)")
    st.write("Paper is [here](https://arxiv.org/abs/2305.14720)")
    
    st.image('content/teaser-website.png')


def main():
    
    with st.sidebar:
        st.title("Navigation")
        st.slider("Guidance scale", 0.0, 20.0, 7.5, 0.1)
        st.slider("Inference steps", 5, 40, 20, 1)
        st.number_input("Seed", 0, 100000, 0, 1)


    tab_names = ["Home", "Edit existing image", "Edit generated image", "Zero-shot generation", "Zero-shot stylization"]
    
    (home_tab,
     edit_existing_image_tab,
     edit_generated_image_tab,
     zero_shot_generation_tab,
     zero_shot_stylization_tab) = st.tabs(tab_names)
    
    with home_tab:
        create_home_tab()
        
    with edit_existing_image_tab:
        create_edit_existing_image_tab()
        
    with edit_generated_image_tab:
        create_edit_generated_image_tab()
        
    with zero_shot_generation_tab:
        create_zero_shot_generation_tab()
        
    with zero_shot_stylization_tab:
        create_zero_shot_stylization_tab()
    
    
if __name__ == "__main__":
    main()