def load_checkpoint(
    model: nn.Module,
    path: str,
    cache_dir: str = None,
    device_map: torch.device | str = "cpu",
    cai_model_name: str = "model",
    strict: bool = False,
    rename_keys: dict = None,  # rename keys in the checkpoint to support fine-tuning with a different model architecture; map old_key_prefix to new_key_prefix
) -> nn.Module:
    """
    Loads a checkpoint into model from a path. Support three types of checkpoints:
        1. huggingface safetensors
        2. local .pt or .pth
        3. colossalai sharded checkpoint

    Args:
        model (nn.Module): The model to load the checkpoint into.
        path (str): The path to the checkpoint.
        cache_dir (str): The directory to cache the downloaded checkpoint.
        device_map (torch.device | str): The device to map the checkpoint to.
        cai_model_name (str): The name of the model in the checkpoint.

    Returns:
        nn.Module: The model with the loaded checkpoint.
    """
    if not os.path.exists(path):
        log_message(f"Checkpoint not found at {path}, trying to download from Hugging Face Hub")
        path = load_from_hf_hub(path, cache_dir)
    assert os.path.exists(path), f"Could not find checkpoint at {path}"

    log_message(f"Loading checkpoint from {path}")
    if path.endswith(".safetensors"):
        ckpt = load_file(path, device='cpu')

        if rename_keys is not None:
            # rename keys in the loaded state_dict with old_key_prefix to with new_key_prefix.
            renamed_ckpt = {}
            for old_key, v in ckpt.items():
                new_key = old_key
                for old_key_prefix, new_key_prefix in rename_keys.items():
                    if old_key_prefix in old_key:
                        new_key = old_key.replace(old_key_prefix, new_key_prefix)
                        print(f"Renamed {old_key} to {new_key} in the loaded state_dict")
                        break
                renamed_ckpt[new_key] = v
            ckpt = renamed_ckpt

        missing, unexpected = model.load_state_dict(ckpt, strict=strict)
        print_load_warning(missing, unexpected)
    elif path.endswith(".pt") or path.endswith(".pth"):
        ckpt = torch.load(path, map_location=device_map)
        missing, unexpected = model.load_state_dict(ckpt, strict=strict)
        print_load_warning(missing, unexpected)
    else:
        assert os.path.isdir(path), f"Invalid checkpoint path: {path}"
        load_from_sharded_state_dict(model, path, model_name=cai_model_name, strict=strict)
    return model

    Loads a checkpoint into model from a path. Support three types of checkpoints:
        1. huggingface safetensors
        2. local .pt or .pth
        3. colossalai sharded checkpoint

    Args:
        model (nn.Module): The model to load the checkpoint into.
        path (str): The path to the checkpoint.
        cache_dir (str): The directory to cache the downloaded checkpoint.
        device_map (torch.device | str): The device to map the checkpoint to.
        cai_model_name (str): The name of the model in the checkpoint.

    Returns:
        nn.Module: The model with the loaded checkpoint.
    
opensora/utils/ckpt.py