diff --git a/comfy/bg_removal_model.py b/comfy/bg_removal_model.py index 7877afd7f..6dec65e63 100644 --- a/comfy/bg_removal_model.py +++ b/comfy/bg_removal_model.py @@ -44,7 +44,14 @@ class BackgroundRemovalModel(): comfy.model_management.load_model_gpu(self.patcher) H, W = image.shape[1], image.shape[2] 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) - out = self.model(pixel_values=pixel_values) + + if pixel_values.shape[0] > 1: + out = torch.cat([ + self.model(pixel_values=pixel_values[i:i+1]) + for i in range(pixel_values.shape[0]) + ], dim=0) + else: + out = self.model(pixel_values=pixel_values) out = torch.nn.functional.interpolate(out, size=(H, W), mode="bicubic", antialias=False) mask = out.sigmoid().to(device=comfy.model_management.intermediate_device(), dtype=comfy.model_management.intermediate_dtype())