import timm
import torch.nn as nn


class TimmBackbone(nn.Module):

    def __init__(self, model_name="vit_base_patch16_224", pretrained=False, **kwargs):
        super().__init__()
        self.feat = timm.create_model(
            model_name,
            pretrained=pretrained,
            num_classes=0,
            **kwargs,
        )
        self.feat_dim = self.feat.num_features

    def forward(self, x):
        return self.feat(x)


def timm_backbone(pretrained=False, **kwargs):
    return TimmBackbone(pretrained=pretrained, **kwargs)