else:
    self.model = MODEL_CLS.from_pretrained(
        model_path, torch_dtype='auto', device_map="auto", attn_implementation='flash_attention_2'
    )

    load_wallaroo = True

    if load_wallaroo:

        load_from = "path/to/checkpoint"

        resume_checkpoint = torch.load(load_from, map_location="cpu")
        new_dict = {}
        for key, value in resume_checkpoint['state_dict'].items():
        
        if 'visual' in key:
            new_dict[key.replace('wallaroo', 'model')] = value

        elif 'model' in key:
            new_dict[key.replace('model', 'language_model').replace('wallaroo', 'model')] = value

        elif 'lm_head' in key:
            new_dict['lm_head.weight'] = value


        m, u = self.model.load_state_dict(new_dict, strict=False)
        del resume_checkpoint

    self.model.eval()

cd scripts/evaluate
sh test_ar_t2i.sh

cd scripts/evaluate
sh test_ar_i2i.sh

@article{Zhu2026Simple,
title   = {# A Simple Baseline for Unifying Understanding, Generation, and Editing via Vanilla Next-token Prediction},
author  = {Jie Zhu, Hanghang Ma, Jia Wang, Yayong Guan, Yanbing Zeng, Lishuai Gao,
Junqiang Wu, Jie Hu, Leye Wang},
journal = {arXiv preprint arXiv:2603.04980},
year    = {2026}
}