Update app.py

app.py

@@ -36,7 +36,7 @@ pipe = StableDiffusionInpaintingPipeline.from_pretrained(
 model = CLIPDensePredT(version='ViT-B/16', reduce_dim=64, complex_trans_conv=True)
 
 model.eval()
-model.load_state_dict(torch.load('./clipseg/weights/rd64-uni.pth', map_location=torch.device('cuda')), strict=False)
+model.load_state_dict(torch.load('./clipseg/weights/rd64-uni.pth', map_location=torch.device(device)), strict=False)
 
 transform = transforms.Compose([
       transforms.ToTensor(),
@@ -46,7 +46,8 @@ transform = transforms.Compose([
 
 def predict(radio, dict, word_mask, prompt=""):
     if(radio == "draw a mask above"):
-        with autocast("cuda"):
+        #with autocast("cuda"):
+        with autocast(enable=(False if device=='cpu' else True)):
             init_image = dict["image"].convert("RGB").resize((512, 512))
             mask = dict["mask"].convert("RGB").resize((512, 512))
     else:
@@ -63,7 +64,8 @@ def predict(radio, dict, word_mask, prompt=""):
         cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
         mask = Image.fromarray(np.uint8(bw_image)).convert('RGB')
         os.remove(filename)
-    with autocast("cuda"):
+    #with autocast("cuda"):
+    with autocast(enable=(False if device=='cpu' else True)):
         images = pipe(prompt = prompt, init_image=init_image, mask_image=mask, strength=0.8)["sample"]
     return images[0]