fix bugs

app.py

@@ -28,8 +28,11 @@ def process_image(image, prompt):
         outputs = model(**inputs)
         preds = outputs.logits
   
-    filename = f"mask.png"
-    plt.imsave(filename, torch.sigmoid(preds))
+    filename = "mask.png"
+    preds = torch.sigmoid(preds)
+    preds[preds >= 0.5] = 1
+    preds[preds < 0.5] = 0
+    plt.imsave(filename, preds)
     return Image.open("mask.png").convert("RGB")
 
 
@@ -39,8 +42,8 @@ def read_content(file_path):
     return content
 
 
-def predict(dict, text_query, reference, scale, seed, step):
-    width, height = dict["image"].size
+def predict(input_image, text_query, reference, scale, seed, step):
+    width, height = input_image.size
     if width < height:
         factor = width / 512.0
         width = 512
@@ -51,10 +54,10 @@ def predict(dict, text_query, reference, scale, seed, step):
         height = 512
         width = int((width / factor) / 8.0) * 8
 
-    init_image = dict["image"].convert("RGB").resize((width, height))
-    mask = dict["mask"].convert("RGB").resize((width, height))
-    print(np.array(mask))
-    print(text_query)
+    init_image = input_image.convert("RGB").resize((width, height))
+    mask = process_image(input_image, text_query).resize((width, height))
+    #mask = dict["mask"].convert("RGB").resize((width, height))
+
     generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None
     output = pipe(
         image=init_image,
@@ -122,7 +125,7 @@ with image_blocks as demo:
         with gr.Box():
             with gr.Row():
                 with gr.Column():
-                    image = gr.Image(source="upload", tool="sketch", elem_id="image_upload", type="pil", label="Source Image")
+                    image = gr.Image(source="upload", elem_id="image_upload", type="pil", label="Source Image")
                     text = gr.Textbox(lines=1, placeholder="Clothing item you want to replace...")
                     reference = gr.Image(source="upload", elem_id="image_upload", type="pil", label="Reference Image")
 
@@ -146,11 +149,25 @@ with image_blocks as demo:
 
             with gr.Row():
                 with gr.Column():
-                    gr.Examples(image_list, inputs=[image],label="Examples - Source Image",examples_per_page=12)
+                    gr.Examples(
+                        image_list, 
+                        inputs=[image], 
+                        label="Examples - Source Image",
+                        examples_per_page=12
+                    )
                 with gr.Column():
-                    gr.Examples(ref_list, inputs=[reference],label="Examples - Reference Image",examples_per_page=12)
+                    gr.Examples(
+                        ref_list, 
+                        inputs=[reference], 
+                        label="Examples - Reference Image",
+                        examples_per_page=12
+                    )
             
-            btn.click(fn=predict, inputs=[image, text, reference, guidance, seed, steps], outputs=[image_out, community_icon, loading_icon, share_button])
+            btn.click(
+                fn=predict, 
+                inputs=[image, text, reference, guidance, seed, steps], 
+                outputs=[image_out, community_icon, loading_icon, share_button]
+            )
             share_button.click(None, [], [], _js=share_js)
 
             gr.HTML(