Run torch models in fp16 under inference_mode

- ViTMatte and BiRefNet load with torch_dtype=float16 (the BiRefNet
  checkpoint ships fp16 anyway; the old .float() upcast is gone) and
  floating-point inputs are cast to half to match.
- torch.no_grad() -> torch.inference_mode() in both predict paths.
- Outputs already downcast via .float() before .numpy(), so downstream
  stays float32. anime-seg is ONNX and unaffected.

Effect verified by the user beforehand.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-07-06 02:50:42 +08:00
parent 5265e802da
commit c7a2819d33
2 changed files with 13 additions and 7 deletions
+5 -4
View File
@@ -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]:
+8 -3
View File
@@ -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()