main.py

import torch
from parameter import get_args
from pretrain import pretrain
from adapt import adapt
from utils.data_utils import *
from backdoor import *
from model.diffusion import *
from preprocess import get_pretrain_model

import warnings; warnings.filterwarnings('ignore')

def main(args):
    model = get_pretrain_model(args)
    
    if args.pretrain:
        pretrain(args, model)
    model.load_state_dict(torch.load(f'{args.save_path}/pretrain.pth', map_location=args.device))
    
    if args.retrain_backdoor:
        if args.gfm_model == "GCOPE":
            gcope_bad_pretrain(args, model)
        elif args.gfm_model == "MDGPT":
            mdgpt_bad_pretrain(args, model)
        elif args.gfm_model == "SAMGPT":
            samgpt_bad_pretrain(args, model)
        else:
            raise NotImplementedError

    if args.backdoor and (args.backdoor_model == "GCBA" or args.backdoor_model == "CrossBA"):
        model.load_state_dict(torch.load(f'{args.save_path}/backdoor.pth', map_location=args.device))

    adapt(args, model)
    
if __name__ == '__main__':
    args = get_args()
    args.save_path = f'./save/{args.gfm_model}/{args.backdoor_model}/{args.name}'
    args.log_file = f'./log/{args.gfm_model}/{args.backdoor_model}/{args.name}{args.index}.txt'
    os.makedirs(os.path.dirname(args.log_file), exist_ok=True)
    main(args)

# python main.py --name {dataset_name} --gfm_model {GFM_model} --backdoor_model ours --device cuda:0