Update ui

app.py

@@ -106,18 +106,19 @@ def remove(gradio_image, rm_guidance_scale=9, num_inference_steps=50, seed=42, s
             AAS_end_layer=END_LAYER, # AAS end layer
             num_inference_steps=num_inference_steps, # number of inference steps # AAS_end_step = int(strength*num_inference_steps)
             generator=generator,
-            guidance_scale=1
+            guidance_scale=1,
+            output_type='pt'
         ).images[0]
         print('Inferece: DONE.')
         
-        # pil_mask = to_pil_image(mask.squeeze(0))
-        # pil_mask_blurred = pil_mask.filter(ImageFilter.GaussianBlur(radius=15))
-        # mask_blurred = to_tensor(pil_mask_blurred).unsqueeze_(0).to(mask.device)
-        # mask_f = 1-(1 - mask) * (1 - mask_blurred)
+        pil_mask = to_pil_image(mask.squeeze(0))
+        pil_mask_blurred = pil_mask.filter(ImageFilter.GaussianBlur(radius=15))
+        mask_blurred = to_tensor(pil_mask_blurred).unsqueeze_(0).to(mask.device)
+        mask_f = 1-(1 - mask) * (1 - mask_blurred)
         
-        # image_1 = image.unsqueeze(0)
+        image_1 = image.unsqueeze(0)
         
-        return source_image, image
+        return source_image_pure, pil_mask, image_1
     except:
         print(traceback.format_exc())
     
@@ -165,11 +166,7 @@ with gr.Blocks() as demo:
                 with gr.Column():
                     run_button = gr.Button("Generate")
 
-            result = ImageSlider(
-                interactive=False,
-                label="Generated Image",
-                type="pil"
-            )
+            result = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery", columns=[3], rows=[1], object_fit="contain", height="auto")
     
     run_button.click(
         fn=remove,