| import torch | |
| from transformers.processing_utils import ProcessorMixin | |
| class MSPProcessor(ProcessorMixin): | |
| attributes = ["feature_extractor", "video_processor", "tokenizer"] | |
| feature_extractor_class = "MSPAudioFeatureExtractor" | |
| video_processor_class = "MSPVisualVideoProcessor" | |
| tokenizer_class = "AutoTokenizer" | |
| def __call__(self, images=None, text=None, videos=None, audio=None, **kwargs): | |
| if getattr(self, "feature_extractor", None) and audio is not None: | |
| if "sampling_rate" in kwargs: | |
| sampling_rate = kwargs["sampling_rate"] | |
| else: | |
| sampling_rate = self.feature_extractor.sampling_rate | |
| if isinstance(audio, (str, bytes)): | |
| audio = self.feature_extractor._load_audio( | |
| audio, sample_rate=sampling_rate | |
| ) | |
| elif ( | |
| isinstance(audio, list) and audio and isinstance(audio[0], (str, bytes)) | |
| ): | |
| audio = [ | |
| self.feature_extractor._load_audio(a, sample_rate=sampling_rate) | |
| for a in audio | |
| ] | |
| if getattr(self, "video_processor", None) and videos is not None: | |
| if isinstance(videos, (str, bytes)): | |
| videos = self.video_processor._load_video(videos) | |
| elif ( | |
| isinstance(videos, list) | |
| and videos | |
| and isinstance(videos[0], (str, bytes)) | |
| ): | |
| videos = [self.video_processor._load_video(v) for v in videos] | |
| if getattr(self, "tokenizer", None) and text is not None: | |
| if isinstance(text, bytes): | |
| text = text.decode("utf-8") | |
| elif isinstance(text, list) and text and isinstance(text[0], bytes): | |
| text = [t.decode("utf-8") for t in text] | |
| outputs = super().__call__( | |
| images=images, text=text, videos=videos, audio=audio, **kwargs | |
| ) | |
| if "device" in kwargs: | |
| device = kwargs["device"] | |
| for key, value in outputs.items(): | |
| if isinstance(value, torch.Tensor): | |
| outputs[key] = value.to(device) | |
| elif ( | |
| isinstance(value, list) | |
| and value | |
| and isinstance(value[0], torch.Tensor) | |
| ): | |
| outputs[key] = [v.to(device) for v in value] | |
| return outputs | |