deepfakemodel.py

from transformers import  PreTrainedModel
from facenet_pytorch import MTCNN, InceptionResnetV1
from .deepfakeconfig import 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
        )


DeepFakeConfig.register_for_auto_class()
DeepFakeModel.register_for_auto_class("AutoModelForImageClassification")