mlx-community/Molmo2-8B-4bit

1"""Video processor class for Molmo2"""2from functools import partial3import os4import warnings5from contextlib import redirect_stdout6from io import BytesIO7from urllib.parse import urlparse8from typing import Optional, Union, Callable9 10import numpy as np11import requests12import einops13import torch14import torchvision.transforms15 16from transformers.image_utils import (17    IMAGENET_STANDARD_MEAN,18    IMAGENET_STANDARD_STD,19    ImageInput,20    PILImageResampling,21    SizeDict,22    validate_kwargs,23)24from transformers.video_utils import (25    VideoInput,26    is_valid_video,27    make_batched_videos,28    make_batched_metadata,29    VideoMetadata,30)31from transformers.processing_utils import Unpack, VideosKwargs32from transformers.video_processing_utils import BaseVideoProcessor33from transformers.utils import logging34from transformers.feature_extraction_utils import BatchFeature35from transformers.utils import (36    is_av_available,37    is_decord_available,38    is_torchcodec_available,39    is_yt_dlp_available,40    TensorType,41    logging,42    to_numpy,43)44 45 46logger = logging.get_logger(__name__)47 48MAX_VIDEO_FPS = 849 50 51def normalize_image(52    image: np.ndarray,53    image_mean: list[float],54    image_std: list[float],55) -> np.ndarray:56    image -= np.array(image_mean, dtype=np.float32)[None, None, :]57    image /= np.array(image_std, dtype=np.float32)[None, None, :]58    return image59 60 61def resize_image(62    image: np.ndarray,63    desired_output_size: list[int],64    resample: PILImageResampling,65) -> np.ndarray:66    if len(image.shape) == 3:67        is_video = False68        image = torch.permute(torch.from_numpy(image), [2, 0, 1])69    else:70        is_video = True71        image = torch.permute(torch.from_numpy(image), [0, 3, 1, 2])72    dtype = image.dtype73    if torch.is_floating_point(image):74        in_min = 0.075        in_max = 1.076        resized = torchvision.transforms.Resize(77            desired_output_size,78            resample,79            antialias=False,80        )(image)81        resized = torch.clip(resized, 0.0, 1.0).to(dtype)82    else:83        assert image.dtype == torch.uint8, "SigLIP expects float images or uint8 images, but got {}".format(image.dtype)84        in_min = 0.085        in_max = 255.086        resized = torchvision.transforms.Resize(87            desired_output_size,88            resample,89            antialias=False,90        )(image)91        resized = torch.clip(resized, 0, 255).to(dtype)92 93    resized = resized.to(torch.float32)94    resized = (resized - in_min) / (in_max - in_min)95 96    if is_video:97        resized = torch.permute(resized, [0, 2, 3, 1]).numpy()98    else:99        resized = torch.permute(resized, [1, 2, 0]).numpy()100 101    return resized102 103 104def build_resized_image(105    image: np.ndarray,106    base_image_input_size: list[int],107    resample: PILImageResampling,108    image_mean: list[float],109    image_std: list[float],110    image_patch_size: int,111) -> tuple[np.ndarray, np.ndarray]:112    resized = resize_image(113        image, base_image_input_size, resample,114    )115    resized = normalize_image(resized, image_mean, image_std)116    if len(resized.shape) == 3:117        resized = np.expand_dims(resized, 0)118    crop_patch_w = base_image_input_size[1] // image_patch_size119    crop_patch_h = base_image_input_size[0] // image_patch_size120    resize_idx = np.arange(crop_patch_w*crop_patch_h).reshape([crop_patch_h, crop_patch_w])121    return resized, resize_idx122 123 124def batch_pixels_to_patches(array: np.ndarray, patch_size: int) -> np.ndarray:125    """Reshape images of [n_images, h, w, 3] -> [n_images, n_patches, pixels_per_patch]"""126    if len(array.shape) == 3:127        n_crops, h, w = array.shape128        h_patches = h//patch_size129        w_patches = w//patch_size130        array = np.reshape(array, [n_crops, h_patches, patch_size, w_patches, patch_size])131        array = np.transpose(array, [0, 1, 3, 2, 4])132        array = np.reshape(array, [n_crops, h_patches*w_patches, patch_size*patch_size])133        return array134    else:135        n_crops, h, w, c = array.shape136        h_patches = h//patch_size137        w_patches = w//patch_size138        array = np.reshape(array, [n_crops, h_patches, patch_size, w_patches, patch_size, c])139        array = np.transpose(array, [0, 1, 3, 2, 4, 5])140        array = np.reshape(array, [n_crops, h_patches*w_patches, patch_size*patch_size*c])141        return array142 143 144def arange_for_pooling(145    idx_arr: np.ndarray,146    pool_h: int,147    pool_w: int,148) -> np.ndarray:149    h_pad = pool_h * ((idx_arr.shape[0] + pool_h - 1) // pool_h) - idx_arr.shape[0]150    w_pad = pool_w * ((idx_arr.shape[1] + pool_w - 1) // pool_w) - idx_arr.shape[1]151    idx_arr = np.pad(idx_arr, [[h_pad//2, (h_pad+1)//2], [w_pad//2, (w_pad+1)//2]],152                     mode='constant',constant_values=-1)153    return einops.rearrange(154        idx_arr, "(h dh) (w dw) -> h w (dh dw)", dh=pool_h, dw=pool_w)155 156 157def image_to_patches_and_grids(158    image: ImageInput,159    base_image_input_size: list[int],160    resample: PILImageResampling,161    image_mean: list[float],162    image_std: list[float],163    image_patch_size: int,164    image_pooling_w: int,165    image_pooling_h: int,166) -> tuple[np.ndarray, np.ndarray, np.ndarray]:167    """168    :return image_grids, the shape of each image after pooling169    :return crops, the image crops to processes with the ViT170    :return pooled_patch_idx, for each patch_id tokens in `image_tokens`, the indices of the171                                patches in `crops` to pool for that token, masked with -1172    """173    if isinstance(base_image_input_size, int):174        base_image_input_size = (base_image_input_size, base_image_input_size)175    176    pooling_w = image_pooling_w177    pooling_h = image_pooling_h178 179    resized, resize_idx = build_resized_image(180        image,181        base_image_input_size,182        resample,183        image_mean,184        image_std,185        image_patch_size,186    )187    pooling_idx = arange_for_pooling(resize_idx, pooling_h, pooling_w)188    h, w = pooling_idx.shape[:2]189    pooling_idx = pooling_idx.reshape([-1, pooling_h*pooling_w])190    image_grid = [h, w]191    return (192        image_grid,193        batch_pixels_to_patches(resized, image_patch_size),194        pooling_idx,195    )196 197 198def get_candidate_target_fps(199    video_fps: Union[int, float],200    sampling_fps: Union[int, float],201    max_fps: Union[int, float] = MAX_VIDEO_FPS,202) -> list[float]:203    """204    Return the subset of `video_fps` factors that remain multiples of `sampling_fps`.205 206    Examples:207        >>> get_candidate_target_fps(video_fps=6, sampling_fps=2)208        [2, 6]209        >>> get_candidate_target_fps(video_fps=5, sampling_fps=1)210        [1, 5]211        >>> get_candidate_target_fps(video_fps=2, sampling_fps=2)212        [2]213        >>> get_candidate_target_fps(video_fps=5, sampling_fps=2)214        Traceback (most recent call last):215            ...216        ValueError: sampling_fps=2 must divide video_fps=5 to produce consistent frame steps.217    """218    video_fps = int(video_fps)219    sampling_fps = int(sampling_fps)220    max_fps = int(max_fps)221 222    if sampling_fps is None:223        raise ValueError("sampling_fps must be provided")224    if video_fps <= 0 or sampling_fps <= 0:225        raise ValueError(f"video_fps and sampling_fps must be positive (got {video_fps}, {sampling_fps})")226    if video_fps % sampling_fps != 0:227        raise ValueError(f"sampling_fps={sampling_fps} must divide video_fps={video_fps}.")228 229    candidates = []230    for candidate in range(sampling_fps, video_fps + 1, sampling_fps):231        if candidate > max_fps:232            break233        if video_fps % candidate == 0:234            candidates.append(float(candidate))235    236    return candidates237 238 239def read_video_decord(240    video_path,241    sample_timestamps_fn: Callable,242    **kwargs,243) -> np.ndarray:244    """245    Decode a video using the Decord backend.246 247    Args:248        video_path (`str`):249            Path to the video file.250        sample_timestamps_fn (`Callable`):251            A callable function that will return timestamps at which the video should be sampled.252 253    Returns:254        tuple[`np.array`, `VideoMetadata`]: A tuple containing:255            - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).256            - `VideoMetadata` object.257    """258    # Lazy import from decord259    import importlib260    decord = importlib.import_module("decord")261 262    vr = decord.VideoReader(uri=video_path, ctx=decord.cpu(0))  # decord has problems with gpu263    video_fps = vr.get_avg_fps()264    total_num_frames = len(vr)265    time_stamps = vr.get_frame_timestamp(list(range(len(vr))))266    duration = time_stamps[-1][1] - time_stamps[0][0]267 268    metadata = VideoMetadata(269        total_num_frames=int(total_num_frames),270        fps=float(video_fps),271        duration=float(duration),272        video_backend="decord",273    )274 275    target_timestamps = sample_timestamps_fn(metadata=metadata, **kwargs)276    target_timestamps = np.array(target_timestamps)277    offset = time_stamps[0, 0]278 279    ix = np.searchsorted(time_stamps[:, 1], target_timestamps + offset, side='right')280    ix = np.minimum(ix, len(time_stamps) - 1)281 282    video = vr.get_batch(ix).asnumpy()283    metadata.update(284        {285            "frames_indices": target_timestamps * video_fps,286            "height": video.shape[1],287            "width": video.shape[2],288        }289    )290    return video, metadata291 292 293def read_video_torchcodec(294    video_path,295    sample_timestamps_fn: Callable,296    **kwargs,297) -> np.ndarray:298    """299    Decode a video using torchcodec decoder.300 301    Args:302        video_path (`str`):303            Path to the video file.304        sample_timestamps_fn (`Callable`):305            A callable function that will return timestamps at which the video should be sampled.306 307    Returns:308        tuple[`np.array`, `VideoMetadata`]: A tuple containing:309            - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).310            - `VideoMetadata` object.311    """312    # Lazy import torchcodec313    import importlib314    torchcodec = importlib.import_module("torchcodec")315 316    decoder = torchcodec.decoders.VideoDecoder(317        video_path,318        # Interestingly `exact` mode takes less than approximate when we load the whole video319        seek_mode="exact",320        # Allow FFmpeg decide on the number of threads for efficiency321        num_ffmpeg_threads=0,322    )323    # If the first frame starts at > 0, we effectively clip the video starting at that time324    # since (most) video players would also skip to that time325    time_offset = decoder.metadata.begin_stream_seconds_from_content326    # Note this duration does assume we started playing at `time_offset`327    duration = decoder.metadata.duration_seconds328 329    metadata = VideoMetadata(330        total_num_frames=decoder.metadata.num_frames,331        fps=decoder.metadata.average_fps,332        duration=duration,333        video_backend="torchcodec",334        height=decoder.metadata.height,335        width=decoder.metadata.width,336    )337 338    target_timestamps = sample_timestamps_fn(metadata=metadata, **kwargs)339 340    # Floating point/rounding issues might cause `target_timestamps` to be very slightly341    # out-of-bounds, to handle this we sanity check then clip them342    assert all(x >= 0 for x in target_timestamps)343    assert all(x < duration+1e-6 for x in target_timestamps)344    # 1e-6 padding since torchcodec can throw out-of-bounds errors even if you ask for the345    # exact boundary value, we should still get the first/last frame anyway346    max_timestamp = decoder.metadata.end_stream_seconds_from_content - 1e-6347    min_timestamp = decoder.metadata.begin_stream_seconds_from_content + 1e-6348    # Note we avoid using numpy ops here to reduce floating precision issues349    timestamps = [x + time_offset for x in target_timestamps]350    timestamps = [max(min_timestamp, min(max_timestamp, x)) for x in timestamps]351 352    video = decoder.get_frames_played_at(timestamps).data.numpy().transpose(0, 2, 3, 1)  # Convert to THWC format353    target_timestamps = np.array(target_timestamps)354    metadata.frames_indices = target_timestamps * metadata.fps355 356    return video, metadata357 358 359def read_video_pyav(360    video_path,361    sample_timestamps_fn: Callable,362    **kwargs,363) -> np.ndarray:364    """365    Decode a video using the PyAV backend.366 367    Args:368        video_path (`str`):369            Path to the video file.370        sample_timestamps_fn (`Callable`):371            A callable function that will return timestamps at which the video should be sampled.372 373    Returns:374        tuple[`np.array`, `VideoMetadata`]: A tuple containing:375            - Numpy array of frames in RGB (shape: [num_frames, height, width, 3]).376            - `VideoMetadata` object.377    """378    # Lazy import torchcodec379    import importlib380    av = importlib.import_module("av")381 382    with av.open(video_path) as container:383        video_stream = container.streams.video[0]384        fps = video_stream.average_rate or video_stream.guessed_rate385        it = container.decode(video=0)386        frames = list(it)387 388        stream = container.streams.video[0]389        start = frames[0].pts * stream.time_base390        container_end = stream.duration391        if container_end is not None:392            container_end *= stream.time_base393        if container_end is None or container_end < frames[-1].pts:394            # Some problem with stream duration, so use the frame PTS directly395            # and guess the duration of the last frame396            end = frames[-1].pts * stream.time_base + 1/fps397        else:398            end = container_end399        duration = float(end - start)400 401        metadata = VideoMetadata(402            total_num_frames=len(frames),403            fps=float(fps),404            duration=float(duration),405            video_backend="pyav",406            height=video_stream.height,407            width=video_stream.width,408        )409 410        target_timestamps = sample_timestamps_fn(metadata=metadata, **kwargs)411        offset = float(start)412 413        target_timestamps = np.array(target_timestamps)414        end_time_stamps = np.array([float(frame.pts * stream.time_base) for frame in frames[1:]] + [duration])415        indices = np.searchsorted(end_time_stamps, target_timestamps + offset, side='right')416        indices = np.minimum(indices, len(end_time_stamps) - 1)417 418        video = np.stack(419            [frames[i].to_ndarray(format="rgb24", channel_last=True) for i in indices],420            axis=0,421        )422 423        metadata.frames_indices = target_timestamps * fps424 425        return video, metadata426 427 428VIDEO_DECODERS = {429    "decord": read_video_decord,430    "torchcodec": read_video_torchcodec,431    "pyav": read_video_pyav,432}433 434 435def load_video(436    video: VideoInput,437    backend: str = "decord",438    sample_timestamps_fn: Optional[Callable] = None,439    **kwargs,440):441    """442    Loads `video` to a numpy array.443 444    Args:445        video (`VideoInput`):446            The video to convert to the numpy array format. Can be a link to video or local path.447        backend (`str`, *optional*, defaults to `"decord"`):448            The backend to use when loading the video. Can be any of ["decord", "pyav", ""torchcodec"]. Defaults to "decord".449        sample_timestamps_fn (`Callable`):450            A callable function that will return timestamps at which the video should be sampled.451    """452 453    # Early exit if provided an array or `PIL` frames454    if not isinstance(video, str):455        metadata = [None] * len(video)456        return video, metadata457 458    if urlparse(video).netloc in ["www.youtube.com", "youtube.com"]:459        if not is_yt_dlp_available():460            raise ImportError("To load a video from YouTube url you have  to install `yt_dlp` first.")461        # Lazy import from yt_dlp462        import importlib463        yt_dlp = importlib.import_module("yt_dlp")464 465        buffer = BytesIO()466        with redirect_stdout(buffer), yt_dlp.YoutubeDL() as f:467            f.download([video])468        bytes_obj = buffer.getvalue()469        file_obj = BytesIO(bytes_obj)470    elif video.startswith("http://") or video.startswith("https://"):471        file_obj = BytesIO(requests.get(video).content)472    elif os.path.isfile(video):473        file_obj = video474    else:475        raise TypeError("Incorrect format used for video. Should be an url linking to an video or a local path.")476 477    # can also load with decord, but not cv2/torchvision478    # both will fail in case of url links479    video_is_url = video.startswith("http://") or video.startswith("https://")480    if video_is_url and backend == "opencv":481        raise ValueError("If you are trying to load a video from URL, you cannot use 'opencv' as backend")482 483    if (484        (not is_decord_available() and backend == "decord")485        or (not is_torchcodec_available() and backend == "torchcodec")486        or (not is_av_available() and backend == "pyav")487    ):488        raise ImportError(489            f"You chose backend={backend} for loading the video but the required library is not found in your environment "490            f"Make sure to install {backend} before loading the video."491        )492    493    video_decoder = VIDEO_DECODERS[backend]494    video, metadata = video_decoder(file_obj, sample_timestamps_fn, **kwargs)495    return video, metadata496 497 498def get_target_fps(499    video_fps: float,500    max_frames: int,501    total_frames: int,502    frame_sample_mode: str,503    candidate_target_fps: tuple[float],504) -> float:505    """506    Get the target fps that best spans the video and has the most frames sampled507    """508    num_frames_sampled = 0509    selected_target_fps = None510    for target_fps in candidate_target_fps:511        step_size = max(int(video_fps / target_fps), 1)512        num_frames_sampled_at_fps = int(total_frames / step_size)513        if num_frames_sampled == 0:514            if "uniform" in frame_sample_mode:515                if num_frames_sampled_at_fps > max_frames:516                    break517            selected_target_fps = target_fps518            num_frames_sampled = num_frames_sampled_at_fps519 520        else:521            # the candidate sampling fps increases so frame count can't decrease522            assert num_frames_sampled <= num_frames_sampled_at_fps523            if num_frames_sampled_at_fps > max_frames:524                # choose the sampling fps that spans the video525                continue526 527            elif num_frames_sampled_at_fps > num_frames_sampled:528                # both are less than max_frames, choose the one with higher density of frames sampled529                selected_target_fps = target_fps530                num_frames_sampled = num_frames_sampled_at_fps531    return selected_target_fps532 533 534def get_frame_times_and_chosen_fps(535    selected_target_fps,536    total_frames,537    max_frames,538    video_fps539):540    if selected_target_fps is None:541        frame_indices = np.linspace(0, total_frames, max_frames, endpoint=False, dtype=int)542    else:543        step_size = max(int(video_fps / selected_target_fps), 1)544        frame_indices = np.arange(0, total_frames, step_size)545    if len(frame_indices) > max_frames:546        frame_indices = frame_indices[:max_frames]547    return selected_target_fps, frame_indices548 549 550class Molmo2VideoProcessorKwargs(VideosKwargs, total=False):551    patch_size: Optional[int]552    pooling_size: Optional[list[int]]553    frame_sample_mode: Optional[str]554    max_fps: Optional[int]555    sampling_fps: Optional[int]556 557 558class Molmo2VideoProcessor(BaseVideoProcessor):559    resample = PILImageResampling.BILINEAR560    size = {"height": 378, "width": 378}561    image_mean = IMAGENET_STANDARD_MEAN562    image_std = IMAGENET_STANDARD_STD563    do_resize = True564    do_rescale = True565    do_normalize = True566    do_convert_rgb = True567    patch_size = 14568    pooling_size = [3, 3]569    do_sample_frames = True570    frame_sample_mode = "uniform_last_frame"571    max_fps = 2572    sampling_fps = 2573    valid_kwargs = Molmo2VideoProcessorKwargs574    model_input_names = ["pixel_values_videos", "video_token_pooling", "video_grids"]575 576    def __init__(self, **kwargs: Unpack[Molmo2VideoProcessorKwargs]):577        super().__init__(**kwargs)578        if self.size is not None and (579            self.size.get("height", None) is None or self.size.get("width", None) is None580        ):581            raise ValueError("size must contain 'height' and 'width' keys.")582 583    def _further_process_kwargs(584        self,585        size: Optional[SizeDict] = None,586        **kwargs,587    ) -> dict:588        """589        Update kwargs that need further processing before being validated590        Can be overridden by subclasses to customize the processing of kwargs.591        """592        if size is not None and ("height" not in size or "width" not in size):593            raise ValueError("size must contain 'height' and 'width' keys.")594 595        return super()._further_process_kwargs(size=size, **kwargs)596 597    def sample_times(598        self,599        metadata: VideoMetadata,600        frame_sample_mode: str,601        num_frames: int,602        max_fps: Optional[int] = None,603        sampling_fps: Optional[int] = None,604        **kwargs,605    ) -> np.ndarray:606        """607        Time-based sampling if an array video is passed608        Args:609            metadata (`VideoMetadata`):610                Metadata of the video containing information about total duration, fps and total number of frames.611            frame_sample_mode (`str`, *optional*):612                Mode to sample frames. Defaults to `self.frame_sample_mode`.613            num_frames (`int`, *optional*):614                Maximum number of frames to sample. Defaults to `self.num_frames`.615            man_fps (`int`, *optional*):616                Maximum frames per second to sample.617            sampling_fps (`int`, *optional*):618                Sampling frames per second. Defaults to `self.sampling_fps`.619                Used when `frame_sample_mode` is `"fps"`.620        """621        frame_sample_mode = frame_sample_mode or self.frame_sample_mode622        num_frames = num_frames or self.num_frames623        sampling_fps = sampling_fps or self.sampling_fps624 625        duration = metadata.duration or metadata.total_num_frames / metadata.fps626        if frame_sample_mode == "fps":627            candidate_target_fps = get_candidate_target_fps(metadata.fps, sampling_fps)628            # Try larger and larger FPSs until we hit one that can't span the video629            target_fps = candidate_target_fps[0]630            for candidate_fps in candidate_target_fps[1:]:631                if num_frames / candidate_fps < duration:632                    break633                target_fps = candidate_fps634            times = np.arange(0, num_frames) / target_fps635            times = times[times < duration]636            return times637        elif frame_sample_mode == "uniform_last_frame":638            if max_fps is not None:639                max_duration = (num_frames-1) / max_fps  # -1 to include the last frame640                if max_duration < duration:641                    times = np.linspace(642                        0, duration, num=num_frames, endpoint=True, dtype=np.float64643                    )644                else:645                    times = np.arange(0.0, stop=duration, step=1/max_fps)646                    times = np.concatenate([times, [duration]], axis=0)647                    assert len(times) <= num_frames648            else:649                times = np.linspace(650                    0, duration, num=num_frames, endpoint=True, dtype=np.float64651                )652            return times653        else:654            raise NotImplementedError(frame_sample_mode)655 656    def sample_frames(657        self,658        metadata: VideoMetadata,659        frame_sample_mode: Optional[str] = None,660        num_frames: Optional[int] = None,661        max_fps: Optional[int] = None,662        sampling_fps: Optional[int] = None,663        **kwargs,664    ) -> np.ndarray:665        """666        Frame-based sampling if an array video is passed667        Args:668            metadata (`VideoMetadata`):669                Metadata of the video containing information about total duration, fps and total number of frames.670            frame_sample_mode (`str`, *optional*):671                Mode to sample frames. Defaults to `self.frame_sample_mode`.672            num_frames (`int`, *optional*):673                Maximum number of frames to sample. Defaults to `self.num_frames`.674            max_fps (`int`, *optional*):675                Maximum frames per second to sample.676            sampling_fps (`int`, *optional*):677                Sampling frames per second. Defaults to `self.sampling_fps`.678                Used when `frame_sample_mode` is `"fps"`.679        """680        frame_sample_mode = frame_sample_mode or self.frame_sample_mode681        num_frames = num_frames or self.num_frames682        sampling_fps = sampling_fps or self.sampling_fps683 684        total_num_frames = metadata.total_num_frames685        if frame_sample_mode == "uniform_last_frame" and max_fps is not None:686            duration = total_num_frames / metadata.fps687            if total_num_frames <= 2:688                return np.arange(total_num_frames).astype(int)689            if duration > (num_frames - 1) / max_fps:  # -1 to include the last frame690                # uniform fallback691                indices = np.linspace(692                    0,693                    total_num_frames - 1,694                    num=min(num_frames, total_num_frames),695                    endpoint=True,696                ).astype(int)697                return indices698            else:699                float_indices = np.arange(700                    0.0, stop=total_num_frames - 1, step=float(metadata.fps / max_fps),701                )702                if np.round(float_indices[-1]) != total_num_frames - 1:703                    float_indices = np.concatenate([float_indices, [total_num_frames - 1]], axis=0)704                indices = np.round(float_indices).astype(int)705                assert indices[-1] < total_num_frames706                assert len(float_indices) <= num_frames707                return indices708        elif frame_sample_mode == "uniform_last_frame":709            indices = np.linspace(710                0, total_num_frames - 1, num=min(num_frames, total_num_frames), endpoint=True,711            ).astype(int)712            return indices713        elif frame_sample_mode == "fps":714            candidate_target_fps = get_candidate_target_fps(metadata.fps, sampling_fps)715            selected_target_fps = get_target_fps(716                metadata.fps,717                num_frames,718                total_num_frames,719                frame_sample_mode,720                candidate_target_fps,721            )722            _, indices = get_frame_times_and_chosen_fps(723                selected_target_fps,724                total_num_frames,725                num_frames,726                metadata.fps,727            )728            return indices729        else:730            raise NotImplementedError(frame_sample_mode)731    732    def fetch_videos(733        self,734        video_url_or_urls: Union[str, list[str], list[list[str]]],735        sample_timestamps_fn=None736    ):737        """738        Convert a single or a list of urls into the corresponding `np.array` objects.739 740        If a single url is passed, the return value will be a single object. If a list is passed a list of objects is741        returned.742        """743        if (744            (not is_decord_available())745            and (not is_torchcodec_available())746            and (not is_av_available())747        ):748            raise ImportError(749                "Molmo2VideoProcessor requires `decord`, `torchcodec`, or `av` to be installed."750            )751 752        if is_decord_available():753            backend = "decord"754        elif is_torchcodec_available():755            warnings.warn(756                "`decord` is not installed and cannot be used to decode the video by default. "757                "Falling back to `torchcodec`."758            )759            backend = "torchcodec"760        else:761            warnings.warn(762                "`decord` is not installed and cannot be used to decode the video by default. "763                "Falling back to `PyAV`."764            )765            backend = "pyav"766 767        if isinstance(video_url_or_urls, list):768            return list(zip(*[self.fetch_videos(x, sample_timestamps_fn=sample_timestamps_fn) for x in video_url_or_urls]))769        else:770            return load_video(video_url_or_urls, backend=backend, sample_timestamps_fn=sample_timestamps_fn)771 772    def _decode_and_sample_videos(773        self,774        videos: VideoInput,775        video_metadata: Union[VideoMetadata, dict],776        do_sample_frames: Optional[bool] = None,777        sample_indices_fn: Optional[Callable] = None,778        sample_timestamps_fn: Optional[Callable] = None,779    ):780        """781        Decode input videos and sample frames if needed.782        """783        videos = make_batched_videos(videos)784        video_metadata = make_batched_metadata(videos, video_metadata=video_metadata)785 786        # Framed-based sampling if an array video is passed787        # Otherwise, time-based sampling with decoding788        if is_valid_video(videos[0]) and do_sample_frames:789            assert video_metadata[0].fps is not None, "FPS must be provided for video input"790            sampled_videos = []791            sampled_metadata = []792            for video, metadata in zip(videos, video_metadata):793                indices = sample_indices_fn(metadata=metadata)794                metadata.frames_indices = indices795                sampled_videos.append(video[indices])796                sampled_metadata.append(metadata)797            videos = sampled_videos798            video_metadata = sampled_metadata799        elif not is_valid_video(videos[0]):800            if sample_indices_fn is None:801                logger.warning(802                    "do_sample_frames is False, but video array is not provided: "803                    "Will decode the video and sample frames using Molmo2's default sampling mode"804                )805            if isinstance(videos[0], list):806                raise ValueError(807                    "A list of images is not supported for video input!"808                )809            else:810                videos, video_metadata = self.fetch_videos(videos, sample_timestamps_fn=sample_timestamps_fn)811        812        return videos, video_metadata813    814    def _prepare_input_videos(815        self,816        videos: VideoInput,817        **kwargs,818    ) -> list[np.ndarray]:819        processed_videos = [to_numpy(video) for video in videos]820        return processed_videos821    822    def preprocess(823        self,824        videos: VideoInput,825        **kwargs: Unpack[Molmo2VideoProcessorKwargs],826    ) -> BatchFeature:827        validate_kwargs(828            captured_kwargs=kwargs.keys(),829            valid_processor_keys=list(self.valid_kwargs.__annotations__.keys()) + ["return_tensors"],830        )831 832        # Set default kwargs from self. This ensures that if a kwarg is not provided833        # by the user, it gets its default value from the instance, or is set to None.834        for kwarg_name in self.valid_kwargs.__annotations__:835            kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None))836        837        do_sample_frames = kwargs.pop("do_sample_frames")838        video_metadata = kwargs.pop("video_metadata")839 840        sample_indices_fn = partial(self.sample_frames, **kwargs) if do_sample_frames else None841        sample_timestamps_fn = partial(self.sample_times, **kwargs)842        videos, video_metadata = self._decode_and_sample_videos(843            videos,844            video_metadata=video_metadata,845            do_sample_frames=do_sample_frames,846            sample_indices_fn=sample_indices_fn,847            sample_timestamps_fn=sample_timestamps_fn,848        )849        videos = self._prepare_input_videos(videos=videos)850 851        kwargs = self._further_process_kwargs(**kwargs)852 853        return_metadata = kwargs.pop("return_metadata")854        preprocessed_videos = self._preprocess(videos=videos, **kwargs)855        if return_metadata:856            preprocessed_videos["video_metadata"] = video_metadata857        return preprocessed_videos858    859    def _preprocess(860        self,861        videos: list[np.ndarray],862        size: Optional[SizeDict] = None,863        resample: Optional[PILImageResampling] = None,864        image_mean: Optional[Union[float, list[float]]] = None,865        image_std: Optional[Union[float, list[float]]] = None,866        do_convert_rgb: Optional[bool] = None,867        patch_size: Optional[int] = None,868        pooling_size: Optional[list[int]] = None,869        return_tensors: Optional[Union[str, TensorType]] = None,870        **kwargs,871    ) -> BatchFeature:872        """873        Preprocess a video for the model.874        Args:875            videos (`VideoInput`):876                Video to preprocess.877            size (`SizeDict`, *optional*, defaults to `self.size`):878                Size of the image after resizing.879            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):880                Resampling filter to use when resizing the image. This can be one of the enum `PILImageResampling`. Only881                has an effect if `do_resize` is set to `True`.882            image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):883                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.884            image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):885                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to886                `True`.887            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):888                Whether to convert the image to RGB.889            patch_size (`int`, *optional*, defaults to `self.patch_size`):890                The spatial patch size of the vision encoder.891            pooling_size (`list[int]`, *optional*, defaults to `self.pooling_size`):892                The pooling size of the vision adapter.893            return_tensors (`str` or `TensorType`, *optional*):894                The type of tensors to return. Can be one of:895                - Unset: Return a list of `np.ndarray`.896                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.897                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.898                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.899                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.900 901        Returns:902            A `BatchFeature` containing the following keys:903                - `pixel_values_videos`: The preprocessed videos.904                - `video_token_pooling`: The indices of the patches in `crops` to pool for each token in `video_tokens`.905                - `video_grids`: The video grids.906        """907        if size.height is None or size.width is None:908            raise ValueError("size must contain 'height' and 'width' keys.")909        910        base_image_input_size = [size.height, size.width]911 912        resample = resample or self.resample913        image_mean = image_mean or self.image_mean914        image_std = image_std or self.image_std915        do_convert_rgb = do_convert_rgb or self.do_convert_rgb916 917        patch_size = patch_size or self.patch_size918        pooling_size = pooling_size or self.pooling_size919 920        image_pooling_h, image_pooling_w = pooling_size921 922        batch_grids = []923        batch_crops = []924        batch_pooled_patches_idx = []925 926        for video in videos:927            all_crops = []928            pooled_patches_idx = []929 930            for frame in video:931                image_grid, crops, pooled_idx = image_to_patches_and_grids(932                    frame,933                    base_image_input_size,934                    resample,935                    image_mean,936                    image_std,937                    patch_size,938                    image_pooling_w,939                    image_pooling_h,940                )941                offset = sum(np.prod(x.shape[:2]) for x in all_crops)942                pooled_idx_with_offset = np.where(pooled_idx >= 0, pooled_idx + offset, pooled_idx)943                pooled_patches_idx.append(pooled_idx_with_offset)944                all_crops.append(crops)945 946            video_grid = np.array([len(video), image_grid[0], image_grid[1]])947            all_crops = np.concatenate(all_crops, 0)948            pooled_patches_idx = np.concatenate(pooled_patches_idx, 0)949 950            batch_grids.append(video_grid)951            batch_crops.append(all_crops)952            batch_pooled_patches_idx.append(pooled_patches_idx)953        954        video_grids = np.stack(batch_grids, 0)955        pixel_values_videos = np.concatenate(batch_crops, 0)956        video_token_pooling = np.concatenate(batch_pooled_patches_idx, 0)957        958        data =dict(959            pixel_values_videos=pixel_values_videos,960            video_token_pooling=video_token_pooling,961            video_grids=video_grids,962        )963 964        return BatchFeature(data, tensor_type=return_tensors)965 966 967Molmo2VideoProcessor.register_for_auto_class()