import streamlit as st
from PIL import Image

import torch
from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
from diffusers import StableDiffusionInpaintPipeline


device = 'cuda' if torch.cuda.is_available() else 'cpu'

seg_processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
seg_model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")


# get Stable Diffusion model for image inpainting
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting",
    torch_dtype=torch.float16,
).to(device)


st.title("Stable Edit - Edit your photos with Stable Diffusion!")

# upload image
filename = st.file_uploader("upload an image")
image = Image.open(filename)
st.image(image)

# Select Area to edit
selection = st.selectbox("Select Area(s) to edit", ("AutoSegment Area", "Draw Custom Area", "Suprise Me"))

# TEMP - DEMO stuff
x = st.slider('Select a value')
st.write(x, 'squared is', x * x)