(
self,
device: Optional[torch.device | str] = None,
dtype: Optional[torch.dtype | str] = None,
attention_implementation: Optional[str] = None,
)
| 21 | self.checkpoint = settings.DETECTOR_MODEL_CHECKPOINT |
| 22 | |
| 23 | def model( |
| 24 | self, |
| 25 | device: Optional[torch.device | str] = None, |
| 26 | dtype: Optional[torch.dtype | str] = None, |
| 27 | attention_implementation: Optional[str] = None, |
| 28 | ) -> EfficientViTForSemanticSegmentation: |
| 29 | if device is None: |
| 30 | device = settings.TORCH_DEVICE_MODEL |
| 31 | if dtype is None: |
| 32 | dtype = settings.MODEL_DTYPE |
| 33 | |
| 34 | config = EfficientViTConfig.from_pretrained(self.checkpoint) |
| 35 | model = EfficientViTForSemanticSegmentation.from_pretrained( |
| 36 | self.checkpoint, |
| 37 | dtype=dtype, |
| 38 | config=config, |
| 39 | ) |
| 40 | model = model.to(device) |
| 41 | model = model.eval() |
| 42 | |
| 43 | logger.debug( |
| 44 | f"Loaded detection model {self.checkpoint} from {EfficientViTForSemanticSegmentation.get_local_path(self.checkpoint)} onto device {device} with dtype {dtype}" |
| 45 | ) |
| 46 | return model |
| 47 | |
| 48 | def processor( |
| 49 | self, |
no test coverage detected