from transformers import  PreTrainedModel
from facenet_pytorch import MTCNN, InceptionResnetV1
from deepfakeconfig import DeepFakeConfig
config = DeepFakeConfig()

class DeepFakeModel(PreTrainedModel):
    config_class = DeepFakeConfig
    def __init__(self, config):
        super().__init__(config)
        self.model = InceptionResnetV1(
            pretrained="vggface2",
            classify=True,
            num_classes=1,
            device=config.DEVICE
        )