Update app.py

app.py

@@ -121,12 +121,13 @@ def predict(radio, dict, word_mask, prompt=""):
         init_image = dict['image'].convert('RGB').resize((imgRes, imgRes))
         filename = f"{uuid.uuid4()}.png"
         plt.imsave(filename,torch.sigmoid(preds[0][0]))
-        img2 = cv2.imread(filename)
-        gray_image = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
-        (thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)
-        cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
-        mask = Image.fromarray(np.uint8(bw_image)).convert('RGB')
-        os.remove(filename)
+        ret,img2 = cv2.imread(filename)
+        if ret == True: 
+            gray_image = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
+            (thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)
+            cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
+            mask = Image.fromarray(np.uint8(bw_image)).convert('RGB')
+            os.remove(filename)
     else:
         img = transform(dict["image"]).unsqueeze(0)
         
@@ -135,25 +136,30 @@ def predict(radio, dict, word_mask, prompt=""):
             img.cuda()
             print ("yes, CUDA is available here !! ")
         
-        model = model.to(torch.device(device))
-        img = img.to(torch.device(device))
-        prompt = labels.to(torch.device(device))
         #------------------
         
         word_masks = [word_mask]
         #with torch.cuda.amp.autocast(): #
         with torch.no_grad():
             preds = model(img.repeat(len(word_masks),1,1,1), word_masks)[0]
+
+        
+        model = model.to(torch.device(device))
+        img = img.to(torch.device(device))
+        prompt = labels.to(torch.device(device))
+        
         init_image = dict['image'].convert('RGB').resize((imgRes, imgRes))
         filename = f"{uuid.uuid4()}.png"
         plt.imsave(filename,torch.sigmoid(preds[0][0]))
-        img2 = cv2.imread(filename)
-        gray_image = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
-        (thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)
-        cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
-        mask = Image.fromarray(np.uint8(bw_image)).convert('RGB')
-        os.remove(filename)
-    with autocast(device): #"cuda"
+        ret,img2 = cv2.imread(filename)
+        
+        if ret == True: 
+            gray_image = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
+            (thresh, bw_image) = cv2.threshold(gray_image, 100, 255, cv2.THRESH_BINARY)
+            cv2.cvtColor(bw_image, cv2.COLOR_BGR2RGB)
+            mask = Image.fromarray(np.uint8(bw_image)).convert('RGB')
+            os.remove(filename)
+        with autocast(device): #"cuda"
         images = pipe(prompt = prompt, init_image=init_image, mask_image=mask, strength=0.8)["sample"]
     return images[0]