Update app.py

app.py

@@ -130,11 +130,11 @@ def predict(radio, dict, word_mask, prompt=""):
         plt.imsave(filename,torch.sigmoid(preds[0][0]))
         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)
+        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)
         
@@ -161,11 +161,12 @@ def predict(radio, dict, word_mask, prompt=""):
         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)
+        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]