OpenMotionLab/MotionGPT

1import torch2 3def load_pretrained(cfg, model, logger, phase="train"):4    logger.info(f"Loading pretrain model from {cfg.TRAIN.PRETRAINED}")5    if phase == "train":6        ckpt_path = cfg.TRAIN.PRETRAINED7    elif phase == "test":8        ckpt_path = cfg.TEST.CHECKPOINTS9        10    state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"]11    model.load_state_dict(state_dict, strict=True)12    return model13 14 15def load_pretrained_vae(cfg, model, logger):16    state_dict = torch.load(cfg.TRAIN.PRETRAINED_VAE,17                            map_location="cpu")['state_dict']18    logger.info(f"Loading pretrain vae from {cfg.TRAIN.PRETRAINED_VAE}")19    # Extract encoder/decoder20    from collections import OrderedDict21    vae_dict = OrderedDict()22    for k, v in state_dict.items():23        if "motion_vae" in k:24            name = k.replace("motion_vae.", "")25            vae_dict[name] = v26        elif "vae" in k:27            name = k.replace("vae.", "")28            vae_dict[name] = v29    if hasattr(model, 'vae'):30        model.vae.load_state_dict(vae_dict, strict=True)31    else:32        model.motion_vae.load_state_dict(vae_dict, strict=True)33    34    return model35