Update app.py

app.py

@@ -73,9 +73,23 @@ transform = transforms.Compose([
 
 def predict(radio, dict, word_mask, prompt=""):
     if(radio == "draw a mask above"):
-        with autocast("cuda"):
+        with autocast(device): #"cuda"
             init_image = dict["image"].convert("RGB").resize((512, 512))
             mask = dict["mask"].convert("RGB").resize((512, 512))
+    else if(radio == "type what to keep"):
+        img = transform(dict["image"]).squeeze(0)
+        word_masks = [word_mask]
+        with torch.no_grad():
+            preds = model(img.repeat(len(word_masks),1,1,1), word_masks)[0]
+        init_image = dict['image'].convert('RGB').resize((512, 512))
+        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)
     else:
         img = transform(dict["image"]).unsqueeze(0)
         word_masks = [word_mask]
@@ -90,7 +104,7 @@ 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(device): #"cuda"
         images = pipe(prompt = prompt, init_image=init_image, mask_image=mask, strength=0.8)["sample"]
     return images[0]
 
@@ -175,7 +189,7 @@ with image_blocks as demo:
         with gr.Column():
             image = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="Upload").style(height=400)
             with gr.Box(elem_id="mask_radio").style(border=False):
-                radio = gr.Radio(["draw a mask above", "type what to mask below"], value="draw a mask above", show_label=False, interactive=True).style(container=False)
+                radio = gr.Radio(["draw a mask above", "type what to mask below", "type what to keep"], value="draw a mask above", show_label=False, interactive=True).style(container=False)
                 word_mask = gr.Textbox(label = "What to find in your image", interactive=False, elem_id="word_mask", placeholder="Disabled").style(container=False)
             prompt = gr.Textbox(label = 'Your prompt (what you want to add in place of what you are removing)')
             radio.change(fn=swap_word_mask, inputs=radio, outputs=word_mask,show_progress=False)
@@ -192,7 +206,7 @@ with image_blocks as demo:
     gr.HTML(
             """
                 <div class="footer">
-                    <p>Model by <a href="https://huggingface.co/CompVis" style="text-decoration: underline;" target="_blank">CompVis</a> and <a href="https://huggingface.co/stabilityai" style="text-decoration: underline;" target="_blank">Stability AI</a> - Inpainting by <a href="https://github.com/nagolinc" style="text-decoration: underline;" target="_blank">nagolinc</a> and <a href="https://github.com/patil-suraj" style="text-decoration: underline;">patil-suraj</a>, inpainting with words by <a href="https://twitter.com/yvrjsharma/" style="text-decoration: underline;" target="_blank">@yvrjsharma</a> and <a href="https://twitter.com/1littlecoder" style="text-decoration: underline;">@1littlecoder</a> - Gradio Demo by 🤗 Hugging Face
+                    <p>Model by <a href="https://huggingface.co/CompVis" style="text-decoration: underline;" target="_blank">CompVis</a> and <a href="https://huggingface.co/stabilityai" style="text-decoration: underline;" target="_blank">Stability AI</a> - Inpainting by <a href="https://github.com/" style="text-decoration: underline;" target="_blank">NightFury</a> using clipseg[model] with bit modification - Gradio Demo on 🤗 Hugging Face
                     </p>
                 </div>
                 <div class="acknowledgments">