Spaces:
Runtime error
Runtime error
File size: 1,624 Bytes
be13417 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
"""
Copyright (c) 2023, salesforce.com, inc.
All rights reserved.
SPDX-License-Identifier: BSD-3-Clause
For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""
from lavis.models.base_model import BaseEncoder
from lavis.models.beats.BEATs import BEATs, BEATsConfig
import torch
from lavis.common.utils import is_url
from lavis.common.dist_utils import download_cached_file
import os
ckp_path = "https://valle.blob.core.windows.net/share/BEATs/BEATs_iter3_plus_AS2M.pt?sv=2020-08-04&st=2023-03-01T07%3A51%3A05Z&se=2033-03-02T07%3A51%3A00Z&sr=c&sp=rl&sig=QJXmSJG9DbMKf48UDIU1MfzIro8HQOf3sqlNXiflY1I%3D"
class BeatsEncoder(BaseEncoder):
def __init__(self, checkpoint_path=ckp_path):
super().__init__()
# load the pre-trained checkpoints
if is_url(checkpoint_path):
cached_file = download_cached_file(
checkpoint_path, check_hash=False, progress=True
)
checkpoint = torch.load(cached_file)
elif os.path.isfile(checkpoint_path):
checkpoint = torch.load(checkpoint_path)
cfg = BEATsConfig(checkpoint['cfg'])
self.num_features = cfg.encoder_embed_dim
self.model = BEATs(cfg)
self.model.load_state_dict(checkpoint['model'])
self.model.eval()
@classmethod
def from_config(cls, cfg):
checkpoint_path = cfg.get("checkpoint_path",ckp_path)
return cls(checkpoint_path)
def forward(self, x):
with torch.no_grad():
return self.model.extract_features(x.squeeze(1))[0]
|