diff --git a/bridge_former.py b/bridge_former.py index 156f2c0..400e6eb 100644 --- a/bridge_former.py +++ b/bridge_former.py @@ -8,7 +8,7 @@ import numpy from towhee import register from towhee.operator.base import NNOperator from towhee.types.video_frame import VideoFrame -from towhee.models.utils.video_transforms import transform_video +from towhee.models.utils.video_transforms import transform_video, get_configs from towhee.models.bridgeformer import bridge_former from transformers import AutoTokenizer @@ -52,7 +52,15 @@ class BridgeFormer(NNOperator): model_name=self.model_name) self.tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased', TOKENIZERS_PARALLELISM=False) - self.transform_cfgs = configs(self.model_name) + # self.transform_cfgs = configs(self.model_name) + self.transform_cfgs = get_configs( + side_size=256, + crop_size=224, + num_frames=None, + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225], + ) + self.model.eval() def decoder_video(self, data: List[VideoFrame]): # Convert list of towhee.types.Image to numpy.ndarray in float32