Update inference_i2mv_sdxl.py

inference_i2mv_sdxl.py

@@ -101,8 +101,8 @@ def remove_bg(image, net, transform, device):
     image_size = image.size
     input_images = transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
-        #preds = net(input_images)[-1].sigmoid().cpu()
-        preds = net(input_images)[-1] if isinstance(net(input_images), list) else net(input_images)
+        preds = net(input_images)[-1][0].sigmoid().cpu()
+        #preds = net(input_images)[-1] if isinstance(net(input_images), list) else net(input_images)
     pred = preds[0].squeeze()
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize(image_size)