Update flow

app.py

@@ -104,8 +104,7 @@ 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,
-        output_type='pt'
+        guidance_scale=1
     ).images[0]
     print('Inferece: DONE.')
     
@@ -114,9 +113,9 @@ def remove(gradio_image, rm_guidance_scale=9, num_inference_steps=50, seed=42, s
     # 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_1
+    return source_image, image
     
 
 title = """<h1 align="center">Object Remove</h1>"""