deepfake / deepfakemodel.py
not-lain's picture
commit files to HF hub
2374ae7
raw
history blame
562 Bytes
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")