diff --git a/bgfilter/segmentation.py b/bgfilter/segmentation.py index f7fa882..5792dcd 100644 --- a/bgfilter/segmentation.py +++ b/bgfilter/segmentation.py @@ -38,10 +38,11 @@ class BiRefNetSegmenter: self.torch = torch self.model = AutoModelForImageSegmentation.from_pretrained( - resolve_model_source(settings.model_name), trust_remote_code=True + resolve_model_source(settings.model_name), + trust_remote_code=True, + torch_dtype=torch.float16, # checkpoint ships fp16; run it as-is, inputs cast to match ) self.model.eval() - self.model.float() # checkpoint ships as fp16; force fp32 to match inputs self.device = self._resolve_device(settings.device) self.model.to(self.device) @@ -64,8 +65,8 @@ class BiRefNetSegmenter: def mask(self, rgb: np.ndarray) -> np.ndarray: """Return a soft foreground mask, float32 0..1, at the input resolution.""" image = Image.fromarray(rgb.astype(np.uint8), mode="RGB") - tensor = self.transform(image).unsqueeze(0).to(self.device) - with self.torch.no_grad(): + tensor = self.transform(image).unsqueeze(0).to(self.device, dtype=self.torch.float16) + with self.torch.inference_mode(): pred = self.model(tensor)[-1].sigmoid() pred = pred[0, 0].detach().float().cpu().numpy() if pred.shape != rgb.shape[:2]: diff --git a/bgfilter/vitmatte_infer.py b/bgfilter/vitmatte_infer.py index 9884f4f..3d5f953 100644 --- a/bgfilter/vitmatte_infer.py +++ b/bgfilter/vitmatte_infer.py @@ -21,7 +21,7 @@ class ViTMatteRunner: self.torch = torch source = resolve_model_source(settings.model_name) self.processor = VitMatteImageProcessor.from_pretrained(source) - self.model = VitMatteForImageMatting.from_pretrained(source) + self.model = VitMatteForImageMatting.from_pretrained(source, torch_dtype=torch.float16) self.device = self._resolve_device(settings.device) self.model.to(self.device) self.model.eval() @@ -37,8 +37,13 @@ class ViTMatteRunner: image = Image.fromarray(rgb.astype(np.uint8), mode="RGB") trimap_image = Image.fromarray(trimap.astype(np.uint8), mode="L") inputs = self.processor(images=image, trimaps=trimap_image, return_tensors="pt") - inputs = {key: value.to(self.device) for key, value in inputs.items()} - with self.torch.no_grad(): + inputs = { + key: value.to(self.device, dtype=self.torch.float16) + if value.is_floating_point() + else value.to(self.device) + for key, value in inputs.items() + } + with self.torch.inference_mode(): outputs = self.model(**inputs) alpha = outputs.alphas[0, 0].detach().float().cpu().numpy()