third_party/VACE/vace/annotators/midas/base_model.py

# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import torch


class BaseModel(torch.nn.Module):
    def load(self, path):
        """Load model from file.

        Args:
            path (str): file path
        """
        parameters = torch.load(path, map_location=torch.device('cpu'), weights_only=True)

        if 'optimizer' in parameters:
            parameters = parameters['model']

        self.load_state_dict(parameters)