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:
@@ -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]:
|
||||
|
||||
Reference in New Issue
Block a user