David310/Detect_AI-generated_Image

1from .clip_models import CLIPModel2from .imagenet_models import ImagenetModel3 4 5VALID_NAMES = [6    'Imagenet:resnet18',7    'Imagenet:resnet34',8    'Imagenet:resnet50',9    'Imagenet:resnet101',10    'Imagenet:resnet152',11    'Imagenet:vgg11',12    'Imagenet:vgg19',13    'Imagenet:swin-b',14    'Imagenet:swin-s',15    'Imagenet:swin-t',16    'Imagenet:vit_b_16',17    'Imagenet:vit_b_32',18    'Imagenet:vit_l_16',19    'Imagenet:vit_l_32',20 21    'CLIP:RN50', 22    'CLIP:RN101', 23    'CLIP:RN50x4', 24    'CLIP:RN50x16', 25    'CLIP:RN50x64', 26    'CLIP:ViT-B/32', 27    'CLIP:ViT-B/16', 28    'CLIP:ViT-L/14', 29    'CLIP:ViT-L/14@336px',30]31 32 33 34 35 36def get_model(name):37    assert name in VALID_NAMES38    if name.startswith("Imagenet:"):39        return ImagenetModel(name[9:]) 40    elif name.startswith("CLIP:"):41        return CLIPModel(name[5:])  42    else:43        assert False 44