Spaces:
Running
on
Zero
Running
on
Zero
Update inference_i2mv_sdxl.py
Browse files- inference_i2mv_sdxl.py +2 -1
inference_i2mv_sdxl.py
CHANGED
@@ -101,7 +101,8 @@ def remove_bg(image, net, transform, device):
|
|
101 |
image_size = image.size
|
102 |
input_images = transform(image).unsqueeze(0).to(device)
|
103 |
with torch.no_grad():
|
104 |
-
preds = net(input_images)[-1].sigmoid().cpu()
|
|
|
105 |
pred = preds[0].squeeze()
|
106 |
pred_pil = transforms.ToPILImage()(pred)
|
107 |
mask = pred_pil.resize(image_size)
|
|
|
101 |
image_size = image.size
|
102 |
input_images = transform(image).unsqueeze(0).to(device)
|
103 |
with torch.no_grad():
|
104 |
+
#preds = net(input_images)[-1].sigmoid().cpu()
|
105 |
+
preds = net(input_images)[-1] if isinstance(net(input_images), list) else net(input_images)
|
106 |
pred = preds[0].squeeze()
|
107 |
pred_pil = transforms.ToPILImage()(pred)
|
108 |
mask = pred_pil.resize(image_size)
|