OmniUE-7B / src /sam_audio /processor.py

first upload

5aee137 17 days ago

17.9 kB

# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved\n

import json
import logging
import math
import os
from typing import Callable, List, Optional, Tuple

import torch
import torchaudio
from huggingface_hub import hf_hub_download
from torch.nn.utils.rnn import pad_sequence
from torchcodec.decoders import AudioDecoder, VideoDecoder
from transformers import AutoTokenizer, BatchFeature

from .model.config import SAMAudioConfig, SAMAudioJudgeConfig

logger = logging.getLogger(__name__)

Anchor = Tuple[str, float, float]


def batch_audio(
    audios: list[str | torch.Tensor], audio_sampling_rate: int = 48_000
) -> Tuple[torch.Tensor, torch.Tensor]:
    wavs = []
    for audio in audios:
        if isinstance(audio, str):
            wav, sr = torchaudio.load(audio)
            if sr != audio_sampling_rate:
                wav = torchaudio.functional.resample(wav, sr, audio_sampling_rate)
        else:
            wav = audio
        wavs.append(wav.mean(0))
    sizes = torch.tensor([wav.size(-1) for wav in wavs])
    return pad_sequence(wavs, batch_first=True).unsqueeze(1), sizes


class Batch:
    def __init__(
        self,
        audios: torch.Tensor,
        sizes: torch.Tensor,
        wav_sizes: torch.Tensor,
        descriptions: list[str],
        hop_length: int,
        audio_sampling_rate: int,
        anchors: Optional[list[list[Anchor]]] = None,
        audio_pad_mask: Optional[torch.Tensor] = None,
        masked_video: Optional[torch.Tensor] = None,
    ):
        self.audios = audios
        self.sizes = sizes
        self.wav_sizes = wav_sizes
        self.descriptions = descriptions
        self.audio_pad_mask = audio_pad_mask
        self.masked_video = masked_video
        self.hop_length = hop_length
        self.audio_sampling_rate = audio_sampling_rate
        self.process_anchors(anchors)
        assert self.audios.size(0) == len(self.descriptions)

    def _wav_to_feature_idx(self, wav_idx: int):
        return math.ceil(wav_idx / self.hop_length)

    def to(self, device: torch.device):
        self.audios = self.audios.to(device)
        self.anchor_ids = self.anchor_ids.to(device)
        self.anchor_alignment = self.anchor_alignment.to(device)
        self.sizes = self.sizes.to(device)
        self.wav_sizes = self.wav_sizes.to(device)
        if self.audio_pad_mask is not None:
            self.audio_pad_mask = self.audio_pad_mask.to(device)
        if self.masked_video is not None:
            self.masked_video = [v.to(device) for v in self.masked_video]
        return self

    def process_anchors(self, anchors: Optional[list[list[Anchor]]]):
        batch_size = len(self.audios)
        anchor_dict = {"<null>": 0, "+": 1, "-": 2, "<pad>": 3}
        if anchors is None:
            anchor_ids = torch.full(
                (batch_size, 2), anchor_dict["<null>"], dtype=torch.long
            )
            anchor_ids[:, 1] = anchor_dict["<pad>"]
            anchor_alignment = torch.full(
                (
                    batch_size,
                    self.audio_pad_mask.size(-1),
                ),
                0,
                dtype=torch.long,
            )
            anchor_alignment[~self.audio_pad_mask] = 1  # point to pad token
        else:
            anchor_alignment = torch.full(
                (
                    batch_size,
                    self.audio_pad_mask.size(-1),
                ),
                0,
                dtype=torch.long,
            )
            anchor_alignment[~self.audio_pad_mask] = 1  # point to pad token
            ids = []

            for i, anchor_list in enumerate(anchors):
                current = [anchor_dict["<null>"], anchor_dict["<pad>"]]
                for token, start_time, end_time in anchor_list:
                    start_idx = self._wav_to_feature_idx(
                        start_time * self.audio_sampling_rate
                    )
                    end_idx = self._wav_to_feature_idx(
                        end_time * self.audio_sampling_rate
                    )
                    anchor_alignment[i, start_idx:end_idx] = len(current)
                    current.append(anchor_dict[token])
                ids.append(torch.tensor(current))
            anchor_ids = pad_sequence(
                ids, batch_first=True, padding_value=anchor_dict["<pad>"]
            )
        self.anchor_ids = anchor_ids
        self.anchor_alignment = anchor_alignment
        self.anchors = anchors


def mask_from_sizes(sizes: torch.Tensor) -> torch.Tensor:
    return torch.arange(sizes.max()).expand(len(sizes), -1) < sizes.unsqueeze(1)


def load_video(
    sizes: torch.Tensor,
    videos: List[str],
    feature_idx_to_wav_idx: Callable[[torch.Tensor], torch.Tensor],
    audio_sampling_rate: int,
) -> list[torch.Tensor]:
    all_frames = []
    for size, video in zip(sizes, videos, strict=False):
        audio_timestamps = (
            feature_idx_to_wav_idx(torch.arange(size)) / audio_sampling_rate
        )
        if isinstance(video, str):
            decoder = VideoDecoder(video, dimension_order="NCHW")
            data = decoder.get_frames_in_range(0, len(decoder))
            diffs = (audio_timestamps[None] - data.pts_seconds[:, None]).abs()
            frame_idxs = diffs.argmin(dim=0)
            frames = data.data[frame_idxs]
        else:
            assert video.size(1) == 3, (
                f"Expected video tensor to be in NCHW format, but found {video.size(1)} channels"
            )
            idx = torch.linspace(0, video.size(0) - 1, int(size)).round().long()
            frames = video[idx]
        all_frames.append(frames)
    return all_frames


class Processor:
    config_cls: Callable

    def __init__(self, audio_hop_length: int, audio_sampling_rate: int):
        self.audio_hop_length = audio_hop_length
        self.audio_sampling_rate = audio_sampling_rate

    @classmethod
    def _get_config(cls, model_name_or_path: str):
        if os.path.exists(model_name_or_path):
            config_path = os.path.join(model_name_or_path, "config.json")
        else:
            config_path = hf_hub_download(
                repo_id=model_name_or_path,
                filename="config.json",
                revision=cls.revision,
            )
        with open(config_path) as fin:
            config = cls.config_cls(**json.load(fin))
        return config

    @classmethod
    def from_pretrained(cls, model_name_or_path: str) -> "Processor":
        config = cls._get_config(model_name_or_path)
        return cls(
            audio_hop_length=config.audio_codec.hop_length,
            audio_sampling_rate=config.audio_codec.sample_rate,
        )

    def feature_to_wav_idx(self, feature_idx):
        return feature_idx * self.audio_hop_length

    def wav_to_feature_idx(self, wav_idx):
        if torch.is_tensor(wav_idx):
            ceil = torch.ceil
        else:
            ceil = math.ceil
        return ceil(wav_idx / self.audio_hop_length)

    @staticmethod
    def _compute_nframes(
        total_frames: int,
        video_fps: float,
        target_fps: float,
        max_frames: Optional[int] = None,
    ) -> int:
        """Compute number of frames to sample — mirrors qwen_omni_utils.smart_nframes."""
        FRAME_FACTOR = 2
        nframes = total_frames / video_fps * target_fps
        nframes = max(nframes, FRAME_FACTOR)
        if max_frames is not None:
            nframes = min(nframes, max_frames)
        nframes = min(nframes, total_frames)
        nframes = int(nframes // FRAME_FACTOR) * FRAME_FACTOR
        nframes = max(nframes, FRAME_FACTOR)
        return nframes

    def mask_videos(
        self,
        videos: List[str | torch.Tensor],
        masks: List[str | torch.Tensor],
        max_frames: Optional[int] = None,
        fps: Optional[float] = None,
    ) -> list[torch.Tensor]:
        video = []
        for v in videos:
            if isinstance(v, str):
                video_path = v
                decoder = VideoDecoder(v, dimension_order="NCHW")
                if fps is not None:
                    # FPS-based selective decoding (same strategy as Qwen-Omni)
                    video_fps = decoder.metadata.average_fps
                    total_frames = decoder.metadata.num_frames
                    nframes = self._compute_nframes(
                        total_frames, video_fps, fps, max_frames
                    )
                    idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
                    decoded = decoder.get_frames_at(indices=idx)
                    v = decoded.data if hasattr(decoded, "data") else decoded
                    logger.debug(
                        "mask_videos: %s — native %d frames @ %.1f fps, "
                        "sampled %d frames (target_fps=%.2f, max_frames=%s)",
                        video_path, total_frames, video_fps, nframes, fps, max_frames,
                    )
                else:
                    decoded = decoder[:]
                    v = decoded.data if hasattr(decoded, "data") else decoded
            elif fps is None and max_frames is not None and v.shape[0] > max_frames:
                idx = torch.linspace(0, v.shape[0] - 1, max_frames).round().long()
                v = v[idx]
            video.append(v)

        video_mask = []
        for m in masks:
            if isinstance(m, str):
                decoded = VideoDecoder(m, dimension_order="NCHW")[:]
                m = decoded.data if hasattr(decoded, "data") else decoded
            video_mask.append(m)

        result = []
        for v, m in zip(video, video_mask, strict=False):
            T_video = v.shape[0]
            T_mask = m.shape[0]

            # Temporal alignment: subsample video to match the mask frame count
            # (masks come from SAM3 which may process fewer frames than the video)
            if T_video > T_mask:
                idx = torch.linspace(0, T_video - 1, T_mask).round().long()
                v = v[idx]
            elif T_mask > T_video:
                idx = torch.linspace(0, T_mask - 1, T_video).round().long()
                m = m[idx]

            # Further subsample both if max_frames is set and fps wasn't used
            # (fps path already applied max_frames during decode)
            if fps is None and max_frames is not None and v.shape[0] > max_frames:
                idx = torch.linspace(0, v.shape[0] - 1, max_frames).round().long()
                v = v[idx]
                m = m[idx]

            # Ensure same device
            m = m.to(device=v.device)
            # Resize mask spatially to match video if needed
            if m.shape[-2:] != v.shape[-2:]:
                orig_shape = m.shape
                flat = m.reshape(-1, 1, *orig_shape[-2:]).float()
                flat = torch.nn.functional.interpolate(
                    flat, size=v.shape[-2:], mode="nearest"
                )
                m = flat.reshape(*orig_shape[:-2], *v.shape[-2:])
            # Union across objects dim so mask is (T, 1, H, W) for broadcasting
            # with video (T, C, H, W)
            if m.dim() == 4 and m.shape[1] != 1 and m.shape[1] != v.shape[1]:
                m = m.amax(dim=1, keepdim=True)
            result.append(v * m.eq(0))
        return result


class SAMAudioProcessor(Processor):
    config_cls = SAMAudioConfig
    revision = None

    def __call__(
        self,
        descriptions: list[str],
        audios: list[str | torch.Tensor],
        anchors: Optional[list[list[Anchor]]] = None,
        masked_videos: Optional[list[str | torch.Tensor]] = None,
    ):
        """
        Processes input data for the model.
        Args:
            descriptions (list[str]): List of text descriptions corresponding to each audio sample.
            audios (list[str]): List of audio file paths or tensors.
                If a tensor:
                    - should have shape (channels, time) where channels=1 for mono and 2 for stereo.
                    - should be resampled to 48_000 hz
            anchors (Optional[list[list[Anchor]]], optional): List of anchors for each sample,
                where each anchor is a tuple (token, start_time, end_time).
            masked_videos (Optional[list[str | torch.Tensor]], optional): List of masked video file paths or tensors.
                If a tensor, should have shape (N, C, H, W)
        Returns:
            Batch: A Batch object containing processed audio, sizes, descriptions, anchor ids, anchor alignment, audio pad mask, and optionally masked video.
        """

        assert len(descriptions) == len(audios)
        assert anchors is None or len(descriptions) == len(anchors)
        assert masked_videos is None or len(descriptions) == len(masked_videos)

        audios, wav_sizes = batch_audio(audios, self.audio_sampling_rate)

        sizes = self.wav_to_feature_idx(wav_sizes)
        audio_pad_mask = mask_from_sizes(sizes)
        masked_video = None
        if masked_videos is not None:
            masked_video = load_video(
                sizes, masked_videos, self.feature_to_wav_idx, self.audio_sampling_rate
            )

        return Batch(
            audios=audios,
            sizes=sizes,
            descriptions=descriptions,
            audio_pad_mask=audio_pad_mask,
            anchors=anchors,
            masked_video=masked_video,
            hop_length=self.audio_hop_length,
            audio_sampling_rate=self.audio_sampling_rate,
            wav_sizes=wav_sizes,
        )


class SAMAudioJudgeProcessor(Processor):
    config_cls = SAMAudioJudgeConfig
    revision = "sam_audio"

    def __init__(
        self,
        audio_hop_length: int,
        audio_sampling_rate: int,
        tokenizer: AutoTokenizer,
    ):
        super().__init__(audio_hop_length, audio_sampling_rate)
        self.tokenizer = tokenizer

    @classmethod
    def from_pretrained(cls, model_name_or_path: str) -> "SAMAudioJudgeProcessor":
        config = cls._get_config(model_name_or_path)
        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
        return cls(
            audio_hop_length=config.audio_codec.hop_length,
            audio_sampling_rate=config.audio_codec.sample_rate,
            tokenizer=tokenizer,
        )

    def _reflect_pad(self, wav):
        if wav.ndim == 1:
            wav = wav.unsqueeze(0)
        if wav.size(-1) % self.audio_hop_length == 0:
            return wav
        p1d = (0, self.audio_hop_length - (wav.size(-1) % self.audio_hop_length))
        return torch.nn.functional.pad(wav, p1d, mode="reflect")

    def _load_audio(self, path: str):
        ad = AudioDecoder(path, sample_rate=self.audio_sampling_rate, num_channels=1)
        return ad.get_all_samples().data

    def _process_audio(
        self,
        raw_audio,
        sampling_rate: Optional[int] = None,
    ):
        from_file = False
        if isinstance(raw_audio, str):
            raw_audio = [raw_audio]

        if isinstance(raw_audio, (list, tuple)) and isinstance(raw_audio[0], str):
            loaded = []
            for audio_file in raw_audio:
                loaded.append(self._load_audio(audio_file))
            raw_audio = loaded
            from_file = True

        if sampling_rate is not None:
            if sampling_rate != self.audio_sampling_rate:
                raise ValueError(
                    f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of"
                    f" {self.audio_sampling_rate}. Please make sure that the provided audio input was sampled with"
                    f" {self.audio_sampling_rate} and not {sampling_rate}."
                )
        elif not from_file:
            logger.warning(
                f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. "
                "Failing to do so can result in silent errors that might be hard to debug."
            )

        if isinstance(raw_audio, list):
            raw_audio = [self._reflect_pad(x).T for x in raw_audio]
        else:
            raw_audio = self._reflect_pad(raw_audio).T

        # verify inputs are valid
        for example in raw_audio:
            if example.ndim > 2:
                raise ValueError(
                    f"Expected input shape (channels, num_samples), but got shape ({example.shape})"
                )

        lengths = torch.tensor([x.size(0) for x in raw_audio])
        input_values = pad_sequence(raw_audio, batch_first=True).transpose(1, 2)
        padding_mask = torch.arange(lengths.max())[None] < lengths[:, None]

        return BatchFeature(
            {"input_values": input_values, "padding_mask": padding_mask}
        )

    def __call__(
        self,
        text: Optional[str] = None,
        input_audio: Optional[
            str | list[str] | torch.Tensor | list[torch.Tensor]
        ] = None,
        separated_audio: Optional[
            str | list[str] | torch.Tensor | list[torch.Tensor]
        ] = None,
        sampling_rate: Optional[int] = None,
        **kwargs,
    ):
        batch = BatchFeature()
        if text is not None:
            batch.update(
                self.tokenizer(
                    text,
                    return_tensors="pt",
                    padding="longest",
                    max_length=512,
                    truncation=True,
                )
            )

        if input_audio is not None:
            batch.update(self._process_audio(input_audio, sampling_rate))

        if separated_audio is not None:
            batch["separated_values"] = self._process_audio(
                separated_audio, sampling_rate
            )["input_values"]

        return batch


__all__ = ["SAMAudioProcessor", "SAMAudioJudgeProcessor", "Batch"]

	# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved\n

	import json
	import logging
	import math
	import os
	from typing import Callable, List, Optional, Tuple

	import torch
	import torchaudio
	from huggingface_hub import hf_hub_download
	from torch.nn.utils.rnn import pad_sequence
	from torchcodec.decoders import AudioDecoder, VideoDecoder
	from transformers import AutoTokenizer, BatchFeature

	from .model.config import SAMAudioConfig, SAMAudioJudgeConfig

	logger = logging.getLogger(__name__)

	Anchor = Tuple[str, float, float]


	def batch_audio(
	audios: list[str \| torch.Tensor], audio_sampling_rate: int = 48_000
	) -> Tuple[torch.Tensor, torch.Tensor]:
	wavs = []
	for audio in audios:
	if isinstance(audio, str):
	wav, sr = torchaudio.load(audio)
	if sr != audio_sampling_rate:
	wav = torchaudio.functional.resample(wav, sr, audio_sampling_rate)
	else:
	wav = audio
	wavs.append(wav.mean(0))
	sizes = torch.tensor([wav.size(-1) for wav in wavs])
	return pad_sequence(wavs, batch_first=True).unsqueeze(1), sizes


	class Batch:
	def __init__(
	self,
	audios: torch.Tensor,
	sizes: torch.Tensor,
	wav_sizes: torch.Tensor,
	descriptions: list[str],
	hop_length: int,
	audio_sampling_rate: int,
	anchors: Optional[list[list[Anchor]]] = None,
	audio_pad_mask: Optional[torch.Tensor] = None,
	masked_video: Optional[torch.Tensor] = None,
	):
	self.audios = audios
	self.sizes = sizes
	self.wav_sizes = wav_sizes
	self.descriptions = descriptions
	self.audio_pad_mask = audio_pad_mask
	self.masked_video = masked_video
	self.hop_length = hop_length
	self.audio_sampling_rate = audio_sampling_rate
	self.process_anchors(anchors)
	assert self.audios.size(0) == len(self.descriptions)

	def _wav_to_feature_idx(self, wav_idx: int):
	return math.ceil(wav_idx / self.hop_length)

	def to(self, device: torch.device):
	self.audios = self.audios.to(device)
	self.anchor_ids = self.anchor_ids.to(device)
	self.anchor_alignment = self.anchor_alignment.to(device)
	self.sizes = self.sizes.to(device)
	self.wav_sizes = self.wav_sizes.to(device)
	if self.audio_pad_mask is not None:
	self.audio_pad_mask = self.audio_pad_mask.to(device)
	if self.masked_video is not None:
	self.masked_video = [v.to(device) for v in self.masked_video]
	return self

	def process_anchors(self, anchors: Optional[list[list[Anchor]]]):
	batch_size = len(self.audios)
	anchor_dict = {"<null>": 0, "+": 1, "-": 2, "<pad>": 3}
	if anchors is None:
	anchor_ids = torch.full(
	(batch_size, 2), anchor_dict["<null>"], dtype=torch.long
	)
	anchor_ids[:, 1] = anchor_dict["<pad>"]
	anchor_alignment = torch.full(
	(
	batch_size,
	self.audio_pad_mask.size(-1),
	),
	0,
	dtype=torch.long,
	)
	anchor_alignment[~self.audio_pad_mask] = 1 # point to pad token
	else:
	anchor_alignment = torch.full(
	(
	batch_size,
	self.audio_pad_mask.size(-1),
	),
	0,
	dtype=torch.long,
	)
	anchor_alignment[~self.audio_pad_mask] = 1 # point to pad token
	ids = []

	for i, anchor_list in enumerate(anchors):
	current = [anchor_dict["<null>"], anchor_dict["<pad>"]]
	for token, start_time, end_time in anchor_list:
	start_idx = self._wav_to_feature_idx(
	start_time * self.audio_sampling_rate
	)
	end_idx = self._wav_to_feature_idx(
	end_time * self.audio_sampling_rate
	)
	anchor_alignment[i, start_idx:end_idx] = len(current)
	current.append(anchor_dict[token])
	ids.append(torch.tensor(current))
	anchor_ids = pad_sequence(
	ids, batch_first=True, padding_value=anchor_dict["<pad>"]
	)
	self.anchor_ids = anchor_ids
	self.anchor_alignment = anchor_alignment
	self.anchors = anchors


	def mask_from_sizes(sizes: torch.Tensor) -> torch.Tensor:
	return torch.arange(sizes.max()).expand(len(sizes), -1) < sizes.unsqueeze(1)


	def load_video(
	sizes: torch.Tensor,
	videos: List[str],
	feature_idx_to_wav_idx: Callable[[torch.Tensor], torch.Tensor],
	audio_sampling_rate: int,
	) -> list[torch.Tensor]:
	all_frames = []
	for size, video in zip(sizes, videos, strict=False):
	audio_timestamps = (
	feature_idx_to_wav_idx(torch.arange(size)) / audio_sampling_rate
	)
	if isinstance(video, str):
	decoder = VideoDecoder(video, dimension_order="NCHW")
	data = decoder.get_frames_in_range(0, len(decoder))
	diffs = (audio_timestamps[None] - data.pts_seconds[:, None]).abs()
	frame_idxs = diffs.argmin(dim=0)
	frames = data.data[frame_idxs]
	else:
	assert video.size(1) == 3, (
	f"Expected video tensor to be in NCHW format, but found {video.size(1)} channels"
	)
	idx = torch.linspace(0, video.size(0) - 1, int(size)).round().long()
	frames = video[idx]
	all_frames.append(frames)
	return all_frames


	class Processor:
	config_cls: Callable

	def __init__(self, audio_hop_length: int, audio_sampling_rate: int):
	self.audio_hop_length = audio_hop_length
	self.audio_sampling_rate = audio_sampling_rate

	@classmethod
	def _get_config(cls, model_name_or_path: str):
	if os.path.exists(model_name_or_path):
	config_path = os.path.join(model_name_or_path, "config.json")
	else:
	config_path = hf_hub_download(
	repo_id=model_name_or_path,
	filename="config.json",
	revision=cls.revision,
	)
	with open(config_path) as fin:
	config = cls.config_cls(**json.load(fin))
	return config

	@classmethod
	def from_pretrained(cls, model_name_or_path: str) -> "Processor":
	config = cls._get_config(model_name_or_path)
	return cls(
	audio_hop_length=config.audio_codec.hop_length,
	audio_sampling_rate=config.audio_codec.sample_rate,
	)

	def feature_to_wav_idx(self, feature_idx):
	return feature_idx * self.audio_hop_length

	def wav_to_feature_idx(self, wav_idx):
	if torch.is_tensor(wav_idx):
	ceil = torch.ceil
	else:
	ceil = math.ceil
	return ceil(wav_idx / self.audio_hop_length)

	@staticmethod
	def _compute_nframes(
	total_frames: int,
	video_fps: float,
	target_fps: float,
	max_frames: Optional[int] = None,
	) -> int:
	"""Compute number of frames to sample — mirrors qwen_omni_utils.smart_nframes."""
	FRAME_FACTOR = 2
	nframes = total_frames / video_fps * target_fps
	nframes = max(nframes, FRAME_FACTOR)
	if max_frames is not None:
	nframes = min(nframes, max_frames)
	nframes = min(nframes, total_frames)
	nframes = int(nframes // FRAME_FACTOR) * FRAME_FACTOR
	nframes = max(nframes, FRAME_FACTOR)
	return nframes

	def mask_videos(
	self,
	videos: List[str \| torch.Tensor],
	masks: List[str \| torch.Tensor],
	max_frames: Optional[int] = None,
	fps: Optional[float] = None,
	) -> list[torch.Tensor]:
	video = []
	for v in videos:
	if isinstance(v, str):
	video_path = v
	decoder = VideoDecoder(v, dimension_order="NCHW")
	if fps is not None:
	# FPS-based selective decoding (same strategy as Qwen-Omni)
	video_fps = decoder.metadata.average_fps
	total_frames = decoder.metadata.num_frames
	nframes = self._compute_nframes(
	total_frames, video_fps, fps, max_frames
	)
	idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
	decoded = decoder.get_frames_at(indices=idx)
	v = decoded.data if hasattr(decoded, "data") else decoded
	logger.debug(
	"mask_videos: %s — native %d frames @ %.1f fps, "
	"sampled %d frames (target_fps=%.2f, max_frames=%s)",
	video_path, total_frames, video_fps, nframes, fps, max_frames,
	)
	else:
	decoded = decoder[:]
	v = decoded.data if hasattr(decoded, "data") else decoded
	elif fps is None and max_frames is not None and v.shape[0] > max_frames:
	idx = torch.linspace(0, v.shape[0] - 1, max_frames).round().long()
	v = v[idx]
	video.append(v)

	video_mask = []
	for m in masks:
	if isinstance(m, str):
	decoded = VideoDecoder(m, dimension_order="NCHW")[:]
	m = decoded.data if hasattr(decoded, "data") else decoded
	video_mask.append(m)

	result = []
	for v, m in zip(video, video_mask, strict=False):
	T_video = v.shape[0]
	T_mask = m.shape[0]

	# Temporal alignment: subsample video to match the mask frame count
	# (masks come from SAM3 which may process fewer frames than the video)
	if T_video > T_mask:
	idx = torch.linspace(0, T_video - 1, T_mask).round().long()
	v = v[idx]
	elif T_mask > T_video:
	idx = torch.linspace(0, T_mask - 1, T_video).round().long()
	m = m[idx]

	# Further subsample both if max_frames is set and fps wasn't used
	# (fps path already applied max_frames during decode)
	if fps is None and max_frames is not None and v.shape[0] > max_frames:
	idx = torch.linspace(0, v.shape[0] - 1, max_frames).round().long()
	v = v[idx]
	m = m[idx]

	# Ensure same device
	m = m.to(device=v.device)
	# Resize mask spatially to match video if needed
	if m.shape[-2:] != v.shape[-2:]:
	orig_shape = m.shape
	flat = m.reshape(-1, 1, *orig_shape[-2:]).float()
	flat = torch.nn.functional.interpolate(
	flat, size=v.shape[-2:], mode="nearest"
	)
	m = flat.reshape(orig_shape[:-2], v.shape[-2:])
	# Union across objects dim so mask is (T, 1, H, W) for broadcasting
	# with video (T, C, H, W)
	if m.dim() == 4 and m.shape[1] != 1 and m.shape[1] != v.shape[1]:
	m = m.amax(dim=1, keepdim=True)
	result.append(v * m.eq(0))
	return result


	class SAMAudioProcessor(Processor):
	config_cls = SAMAudioConfig
	revision = None

	def __call__(
	self,
	descriptions: list[str],
	audios: list[str \| torch.Tensor],
	anchors: Optional[list[list[Anchor]]] = None,
	masked_videos: Optional[list[str \| torch.Tensor]] = None,
	):
	"""
	Processes input data for the model.

	Args:
	descriptions (list[str]): List of text descriptions corresponding to each audio sample.
	audios (list[str]): List of audio file paths or tensors.
	If a tensor:
	- should have shape (channels, time) where channels=1 for mono and 2 for stereo.
	- should be resampled to 48_000 hz
	anchors (Optional[list[list[Anchor]]], optional): List of anchors for each sample,
	where each anchor is a tuple (token, start_time, end_time).
	masked_videos (Optional[list[str \| torch.Tensor]], optional): List of masked video file paths or tensors.
	If a tensor, should have shape (N, C, H, W)

	Returns:
	Batch: A Batch object containing processed audio, sizes, descriptions, anchor ids, anchor alignment, audio pad mask, and optionally masked video.
	"""

	assert len(descriptions) == len(audios)
	assert anchors is None or len(descriptions) == len(anchors)
	assert masked_videos is None or len(descriptions) == len(masked_videos)

	audios, wav_sizes = batch_audio(audios, self.audio_sampling_rate)

	sizes = self.wav_to_feature_idx(wav_sizes)
	audio_pad_mask = mask_from_sizes(sizes)
	masked_video = None
	if masked_videos is not None:
	masked_video = load_video(
	sizes, masked_videos, self.feature_to_wav_idx, self.audio_sampling_rate
	)

	return Batch(
	audios=audios,
	sizes=sizes,
	descriptions=descriptions,
	audio_pad_mask=audio_pad_mask,
	anchors=anchors,
	masked_video=masked_video,
	hop_length=self.audio_hop_length,
	audio_sampling_rate=self.audio_sampling_rate,
	wav_sizes=wav_sizes,
	)


	class SAMAudioJudgeProcessor(Processor):
	config_cls = SAMAudioJudgeConfig
	revision = "sam_audio"

	def __init__(
	self,
	audio_hop_length: int,
	audio_sampling_rate: int,
	tokenizer: AutoTokenizer,
	):
	super().__init__(audio_hop_length, audio_sampling_rate)
	self.tokenizer = tokenizer

	@classmethod
	def from_pretrained(cls, model_name_or_path: str) -> "SAMAudioJudgeProcessor":
	config = cls._get_config(model_name_or_path)
	tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
	return cls(
	audio_hop_length=config.audio_codec.hop_length,
	audio_sampling_rate=config.audio_codec.sample_rate,
	tokenizer=tokenizer,
	)

	def _reflect_pad(self, wav):
	if wav.ndim == 1:
	wav = wav.unsqueeze(0)
	if wav.size(-1) % self.audio_hop_length == 0:
	return wav
	p1d = (0, self.audio_hop_length - (wav.size(-1) % self.audio_hop_length))
	return torch.nn.functional.pad(wav, p1d, mode="reflect")

	def _load_audio(self, path: str):
	ad = AudioDecoder(path, sample_rate=self.audio_sampling_rate, num_channels=1)
	return ad.get_all_samples().data

	def _process_audio(
	self,
	raw_audio,
	sampling_rate: Optional[int] = None,
	):
	from_file = False
	if isinstance(raw_audio, str):
	raw_audio = [raw_audio]

	if isinstance(raw_audio, (list, tuple)) and isinstance(raw_audio[0], str):
	loaded = []
	for audio_file in raw_audio:
	loaded.append(self._load_audio(audio_file))
	raw_audio = loaded
	from_file = True

	if sampling_rate is not None:
	if sampling_rate != self.audio_sampling_rate:
	raise ValueError(
	f"The model corresponding to this feature extractor: {self} was trained using a sampling rate of"
	f" {self.audio_sampling_rate}. Please make sure that the provided audio input was sampled with"
	f" {self.audio_sampling_rate} and not {sampling_rate}."
	)
	elif not from_file:
	logger.warning(
	f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. "
	"Failing to do so can result in silent errors that might be hard to debug."
	)

	if isinstance(raw_audio, list):
	raw_audio = [self._reflect_pad(x).T for x in raw_audio]
	else:
	raw_audio = self._reflect_pad(raw_audio).T

	# verify inputs are valid
	for example in raw_audio:
	if example.ndim > 2:
	raise ValueError(
	f"Expected input shape (channels, num_samples), but got shape ({example.shape})"
	)

	lengths = torch.tensor([x.size(0) for x in raw_audio])
	input_values = pad_sequence(raw_audio, batch_first=True).transpose(1, 2)
	padding_mask = torch.arange(lengths.max())[None] < lengths[:, None]

	return BatchFeature(
	{"input_values": input_values, "padding_mask": padding_mask}
	)

	def __call__(
	self,
	text: Optional[str] = None,
	input_audio: Optional[
	str \| list[str] \| torch.Tensor \| list[torch.Tensor]
	] = None,
	separated_audio: Optional[
	str \| list[str] \| torch.Tensor \| list[torch.Tensor]
	] = None,
	sampling_rate: Optional[int] = None,
	**kwargs,
	):
	batch = BatchFeature()
	if text is not None:
	batch.update(
	self.tokenizer(
	text,
	return_tensors="pt",
	padding="longest",
	max_length=512,
	truncation=True,
	)
	)

	if input_audio is not None:
	batch.update(self._process_audio(input_audio, sampling_rate))

	if separated_audio is not None:
	batch["separated_values"] = self._process_audio(
	separated_audio, sampling_rate
	)["input_values"]

	return batch


	__all__ = ["SAMAudioProcessor", "SAMAudioJudgeProcessor", "Batch"]