Derive model dtype from device: fp16 on CUDA only, fp32 on CPU

Benchmarking c7a2819 on CPU showed unconditional fp16 is pathological
there: ~1870 s/image vs ~14.4 s fp32 (~130x) -- x86 has no native fp16
conv/gemm kernels so PyTorch falls back to a slow reference path. fp16
only pays off on CUDA, so ViTMatteRunner and BiRefNetSegmenter now pick
float16 iff the resolved device is cuda, float32 otherwise; inputs cast
to the same dtype. The dtype is intentionally not a setting.

Also replaces the deprecated torch_dtype= kwarg with dtype=.

Verified on CPU: single-image CLI run back to normal speed (~31 s wall
including model load) and output byte-identical (SHA256) to the
fp32-era baseline, confirming inference_mode alone changes nothing.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-07-06 11:08:29 +08:00
parent c7a2819d33
commit 291e9faac5
2 changed files with 12 additions and 6 deletions
+6 -3
View File
@@ -37,13 +37,16 @@ class BiRefNetSegmenter:
) from exc
self.torch = torch
self.device = self._resolve_device(settings.device)
# fp16 only pays off on CUDA; CPU lacks native fp16 kernels (~100x slower),
# so the dtype follows the device. The fp16 checkpoint upcasts on CPU.
self.dtype = torch.float16 if self.device.type == "cuda" else torch.float32
self.model = AutoModelForImageSegmentation.from_pretrained(
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
dtype=self.dtype,
)
self.model.eval()
self.device = self._resolve_device(settings.device)
self.model.to(self.device)
size = settings.input_size
@@ -65,7 +68,7 @@ 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, dtype=self.torch.float16)
tensor = self.transform(image).unsqueeze(0).to(self.device, dtype=self.dtype)
with self.torch.inference_mode():
pred = self.model(tensor)[-1].sigmoid()
pred = pred[0, 0].detach().float().cpu().numpy()
+6 -3
View File
@@ -19,10 +19,13 @@ class ViTMatteRunner:
) from exc
self.torch = torch
self.device = self._resolve_device(settings.device)
# fp16 only pays off on CUDA (tensor cores); CPU has no native fp16 conv/gemm
# kernels and degrades ~100x, so the dtype follows the device, never a param.
self.dtype = torch.float16 if self.device.type == "cuda" else torch.float32
source = resolve_model_source(settings.model_name)
self.processor = VitMatteImageProcessor.from_pretrained(source)
self.model = VitMatteForImageMatting.from_pretrained(source, torch_dtype=torch.float16)
self.device = self._resolve_device(settings.device)
self.model = VitMatteForImageMatting.from_pretrained(source, dtype=self.dtype)
self.model.to(self.device)
self.model.eval()
@@ -38,7 +41,7 @@ class ViTMatteRunner:
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, dtype=self.torch.float16)
key: value.to(self.device, dtype=self.dtype)
if value.is_floating_point()
else value.to(self.device)
for key, value in inputs.items()