update

app.py

@@ -6,18 +6,12 @@ from collections import defaultdict
 import streamlit as st
 from streamlit_drawable_canvas import st_canvas
 
-import torch
-from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
-from diffusers import StableDiffusionInpaintPipeline
-
 import matplotlib as mpl
 
-from model import segment_image, inpaint
+from model import device, segment_image, inpaint
 
 
 # define utils and helpers
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
-
 def closest_number(n, m=8):
     """ Obtains closest number to n that is divisble by m """
     return int(n/m) * m
@@ -96,7 +90,10 @@ def get_mask(image, edit_method, height, width):
 
 if __name__ == '__main__':
     
-    st.title("Stable Edit - Edit your photos with Stable Diffusion!")
+    st.title("Stable Edit")
+    st.title("Edit your photos with Stable Diffusion!")
+
+    st.write(f"Device found: {device}")
 
     sf = st.text_input("Please enter resizing scale factor to downsize image (default=2)", value="2")
     try: 
@@ -126,22 +123,16 @@ if __name__ == '__main__':
 
             # get inpainted images
             prompt = st.text_input("Please enter prompt for image inpainting", value="")
-            seed = st.text_input("(Optional) enter seed to change inpainting result (default=0)", value="0")
-            try:
-                seed = int(seed)
-            except:
-                st.write("Invalid seed! Defaultign to 0, please re-enter above to change it")
-                seed = 0
-
-            st.write("Inpainting Images, patience is a virtue :)")
-
-            images = inpaint(image, mask, width, height, prompt=prompt, seed=seed, guidance_scale=17.5, num_samples=3)
-
-            # display all images
-            st.write("Original Image")
-            st.image(image)
-            for i, img in enumerate(images, 1):
-                st.write(f"result: {i}")
-                st.image(img)
+
+            if prompt: #  and isinstance(seed, int):
+                st.write("Inpainting Images, patience is a virtue and this will take a while to run on a CPU :)")
+                images = inpaint(image, mask, width, height, prompt=prompt, seed=0, guidance_scale=17.5, num_samples=3)
+
+                # display all images
+                st.write("Original Image")
+                st.image(image)
+                for i, img in enumerate(images, 1):
+                    st.write(f"result: {i}")
+                    st.image(img)
 
     

model.py

@@ -1,6 +1,6 @@
 import torch
 from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
-from diffusers import StableDiffusionInpaintPipeline
+from diffusers import StableDiffusionInpaintPipeline # , DiffusionPipeline
 
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -30,7 +30,7 @@ def segment_image(image):
 
     return seg_prediction, segment_labels
 
-# Image inpainting
+# Image inpainting pipeline
 # get Stable Diffusion model for image inpainting
 if device == 'cuda':
     pipe = StableDiffusionInpaintPipeline.from_pretrained(
@@ -39,9 +39,18 @@ if device == 'cuda':
     ).to(device)
 else:
     pipe = StableDiffusionInpaintPipeline.from_pretrained(
-        "runwayml/stable-diffusion-inpainting"
-    ).to(device)
-
+        "runwayml/stable-diffusion-inpainting",
+        torch_dtype=torch.bfloat16, 
+        device_map="auto"
+    )
+
+# pipe = StableDiffusionInpaintPipeline.from_pretrained( # DiffusionPipeline.from_pretrained(
+#         "runwayml/stable-diffusion-inpainting",
+#         revision="fp16",
+#         torch_dtype=torch.bfloat16, 
+#         # device_map="auto" # use for Hugging face spaces
+#     )
+# pipe.to(device) # use for local environment
 
 def inpaint(image, mask, W, H, prompt="", seed=0, guidance_scale=17.5, num_samples=3):
     """ Uses Stable Diffusion model to inpaint image