import streamlit as st
from PIL import Image
import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image

pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16")
generator = torch.Generator().manual_seed(0)

def main():
    st.title("Image and Text Preview App")

    # Upload images
    st.sidebar.subheader("Upload Images")
    uploaded_image1 = st.sidebarfile_uploader("Upload Image", type=["jpg", "jpeg", "png"])
    uploaded_image2 = st.sidebarfile_uploader("Upload Mask", type=["jpg", "jpeg", "png"])

    # Input text
    st.subheader("Text Prompts")
    prompt = st.text_input("Text prompt", "")
    negative_prompt = st.text_input("Negative prompt", "")

    guidance_scale = st.text_input("Guidance scale", 7.5)
    steps = st.text_input("Steps", 20)
    strength = st.text_input("Strength", 1)
    
    btn_submit = st.button("Generate")
     
    if btn_submit:
        init_image = Image.open(uploaded_image1).convert("RGB")
        st.image(init_image, caption="Uploaded Image", use_column_width=True)
        mask_image = Image.open(uploaded_image2).convert("RGB")
        st.image(mask_image, caption="Uploaded Mask", use_column_width=True)
    
        output = pipe(prompt = prompt, 
                      negative_prompt=negative_prompt, 
                      image=init_image, 
                      mask_image=mask_image, 
                      guidance_scale=int(guidance_scale), 
                      num_inference_steps=int(steps), 
                      strength=strength,
                      generator=generator)
        
        st.text("Generated Image:")
        st.image(output.images[0])


if __name__ == "__main__":
    main()