Spaces:
Running
Running
| import functools | |
| from typing import Mapping | |
| import torch | |
| import torch.nn as nn | |
| from networks.base_model import BaseModel | |
| import sys | |
| from models import get_model | |
| class Validator(BaseModel): | |
| def name(self): | |
| return 'Validator' | |
| def __init__(self, opt): | |
| super(Validator, self).__init__(opt) | |
| self.opt = opt | |
| self.model = get_model("FeatureTransformer") | |
| self.clip_model = get_model("CLIP:ViT-L/14") | |
| # for name, p in self.clip_model.named_parameters(): | |
| # if name=="fc.weight" or name=="fc.bias": | |
| # params.append(p) | |
| # else: | |
| # p.requires_grad = False | |
| # del params | |
| self.model.to(self.device) | |
| def set_input(self, input): | |
| # self.input = torch.cat([self.clip_model.forward(x=video_frames, return_feature=True).unsqueeze(0) for video_frames in input[0]]) | |
| self.clip_model.to(self.device) | |
| self.input = self.clip_model.forward(x=input[0].to(self.device).view(-1, 3, 224, 224), return_feature=True).view(-1, 16, 768) | |
| self.clip_model.to('cpu') | |
| self.input = self.input.to(self.device) | |
| self.label = input[1].to(self.device).float() | |
| def forward(self): | |
| self.output = self.model(self.input) | |
| self.output = self.output.view(-1).unsqueeze(1) | |
| def load_state_dict(self, ckpt_path): | |
| state_dict = torch.load(ckpt_path, map_location='cpu') | |
| self.model.load_state_dict(state_dict['model']) | |
| self.model.eval() |