update

app.py

@@ -1,22 +1,22 @@
 import streamlit as st
 from PIL import Image
 
-# import torch
-# from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
-# from diffusers import StableDiffusionInpaintPipeline
+import torch
+from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
+from diffusers import StableDiffusionInpaintPipeline
 
 
-# device = 'cuda' if torch.cuda.is_available() else 'cpu'
+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")
+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)
+# 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!")
@@ -27,7 +27,7 @@ image = Image.open(filename)
 st.image(image)
 
 # Select Area to edit
-st.selectbox("Options", ("AutoSegment Area", "Draw Custom Area"))
+selection = st.selectbox("Select Area(s) to edit", ("AutoSegment Area", "Draw Custom Area", "Suprise Me"))
 
 # TEMP - DEMO stuff
 x = st.slider('Select a value')