Skip to content

Commit 187e523

Browse files
authored
Fix BiRefNet issue (Comfy-Org#13966)
1 parent 164a9d4 commit 187e523

1 file changed

Lines changed: 8 additions & 1 deletion

File tree

comfy/bg_removal_model.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,14 @@ def encode_image(self, image):
4444
comfy.model_management.load_model_gpu(self.patcher)
4545
H, W = image.shape[1], image.shape[2]
4646
pixel_values = comfy.clip_model.clip_preprocess(image.to(self.load_device), size=self.image_size, mean=self.image_mean, std=self.image_std, crop=False)
47-
out = self.model(pixel_values=pixel_values)
47+
48+
if pixel_values.shape[0] > 1:
49+
out = torch.cat([
50+
self.model(pixel_values=pixel_values[i:i+1])
51+
for i in range(pixel_values.shape[0])
52+
], dim=0)
53+
else:
54+
out = self.model(pixel_values=pixel_values)
4855
out = torch.nn.functional.interpolate(out, size=(H, W), mode="bicubic", antialias=False)
4956

5057
mask = out.sigmoid().to(device=comfy.model_management.intermediate_device(), dtype=comfy.model_management.intermediate_dtype())

0 commit comments

Comments
 (0)