Update app.py

app.py

@@ -58,7 +58,7 @@ from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 ##########################
 
 sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
-sam.to(device="cuda")
+#sam.to(device="cpu")
 
 import streamlit as st
 import torch
@@ -77,7 +77,7 @@ import copy
 def load_sam():
     sam_checkpoint = "sam_vit_h_4b8939.pth"
     model_type = "vit_h"
-    device = "cuda"
+    device = "cpu"
 
     sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
     sam.to(device=device)
@@ -104,7 +104,7 @@ def load_pipeline():
         torch_dtype=torch.float16,
         revision="fp16"
     )
-    pipe = pipe.to('cuda')
+    #pipe = pipe.to('cuda')
     pipe.enable_xformers_memory_efficient_attention()
     return pipe
 
@@ -179,7 +179,7 @@ if uploaded_file:
         stable_diffusion_mask = Image.fromarray(selected_mask * 255).convert("RGB")
 
         pipe = load_pipeline()
-        generator = torch.Generator(device="cuda").manual_seed(77)
+        generator = torch.Generator(device="cpu").manual_seed(77)
 
         num_images = 4
         images = []