646b85544b · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -3,7 +3,7 @@

3 import math

4 from collections.abc import Iterable, Mapping, Sequence

5 from dataclasses import dataclass, fields

6 -from functools import cached_property, partial

6 +from functools import partial

7 from itertools import islice

8 from typing import Annotated, Any

9

@@ -14,14 +14,14 @@ import torch.nn.functional as F

14 from PIL import ImageOps

15 from PIL.Image import Image

16 from transformers import (

17 + BaseImageProcessor,

18 + BaseVideoProcessor,

19 BatchFeature,

20 PretrainedConfig,

21 ProcessorMixin,

20 - TensorType,

22 )

23 from transformers.image_utils import ImageInput

23 -from transformers.tokenization_utils_base import TextInput

24 -from transformers.video_utils import VideoInput, VideoMetadata

24 +from transformers.video_utils import VideoMetadata

25

26 from vllm.compilation.decorators import support_torch_compile

27 from vllm.config import CacheConfig, VllmConfig

@@ -1337,12 +1337,14 @@ def exif_transpose(

1337

1338 def build_flat_image_bool_length(

1339 image_grids: torch.LongTensor,

1340 - image_patch_id: int,

1341 - low_res_image_start_id: int,

1342 - image_start_id: int,

1343 - image_col_id: int,

1344 - image_end_id: int,

1340 + hf_config: PretrainedConfig,

1341 ) -> tuple[torch.LongTensor, torch.LongTensor]:

1342 + image_patch_id = hf_config.image_patch_id

1343 + low_res_image_start_id = hf_config.low_res_image_start_token_id

1344 + image_start_id = hf_config.image_start_token_id

1345 + image_col_id = hf_config.image_col_id

1346 + image_end_id = hf_config.image_end_token_id

1347 +

1348 device = image_grids.device

1349 B = image_grids.shape[0]

1350

@@ -1401,10 +1403,12 @@ def build_flat_image_bool_length(

1403

1404 def build_flat_video_bool_length(

1405 video_grids: torch.LongTensor,

1404 - image_patch_id: int,

1405 - frame_start_id: int,

1406 - frame_end_id: int,

1406 + hf_config: PretrainedConfig,

1407 ) -> tuple[torch.LongTensor, torch.LongTensor]:

1408 + image_patch_id = hf_config.image_patch_id

1409 + frame_start_id = hf_config.frame_start_token_id

1410 + frame_end_id = hf_config.frame_end_token_id

1411 +

1412 device = video_grids.device

1413 B = video_grids.shape[0]

1414

@@ -1439,314 +1443,6 @@ def build_flat_video_bool_length(

1443 return flat, lengths

1444

1445

1442 -class Molmo2ProcessorWrapper:

1443 - """

1444 - Wraps :class:`Molmo2Processor` so that it can be called directly.

1445 - """

1446 -

1447 - def __init__(self, processor: ProcessorMixin, hf_config: PretrainedConfig):

1448 - super().__init__()

1449 -

1450 - self.processor = processor

1451 - self.hf_config = hf_config

1452 -

1453 - @cached_property

1454 - def vocab(self) -> dict[str, int]:

1455 - return self.processor.tokenizer.vocab # type: ignore

1456 -

1457 - @cached_property

1458 - def max_crops(self) -> int:

1459 - image_processor = self.processor.image_processor # type: ignore

1460 -

1461 - max_crops = image_processor.max_crops

1462 - assert isinstance(max_crops, int)

1463 -

1464 - return max_crops

1465 -

1466 - @cached_property

1467 - def image_pooling_h(self) -> int:

1468 - image_processor = self.processor.image_processor # type: ignore

1469 -

1470 - image_pooling_h = image_processor.pooling_size[0]

1471 - assert isinstance(image_pooling_h, int)

1472 -

1473 - return image_pooling_h

1474 -

1475 - @cached_property

1476 - def image_pooling_w(self) -> int:

1477 - image_processor = self.processor.image_processor # type: ignore

1478 -

1479 - image_pooling_w = image_processor.pooling_size[1]

1480 - assert isinstance(image_pooling_w, int)

1481 -

1482 - return image_pooling_w

1483 -

1484 - @cached_property

1485 - def video_pooling_h(self) -> int:

1486 - video_processor = self.processor.video_processor # type: ignore

1487 -

1488 - video_pooling_h = video_processor.pooling_size[0]

1489 - assert isinstance(video_pooling_h, int)

1490 -

1491 - return video_pooling_h

1492 -

1493 - @cached_property

1494 - def video_pooling_w(self) -> int:

1495 - video_processor = self.processor.video_processor # type: ignore

1496 -

1497 - video_pooling_w = video_processor.pooling_size[1]

1498 - assert isinstance(video_pooling_w, int)

1499 -

1500 - return video_pooling_w

1501 -

1502 - @cached_property

1503 - def base_image_input_size(self) -> tuple[int, int]:

1504 - if getattr(self.processor, "image_processor", None) is not None:

1505 - processor = self.processor.image_processor # type: ignore

1506 - else:

1507 - processor = self.processor.video_processor # type: ignore

1508 -

1509 - base_image_input_size = (processor.size["height"], processor.size["width"])

1510 -

1511 - return base_image_input_size

1512 -

1513 - @cached_property

1514 - def image_patch_size(self) -> int:

1515 - if getattr(self.processor, "image_processor", None) is not None:

1516 - processor = self.processor.image_processor # type: ignore

1517 - else:

1518 - processor = self.processor.video_processor # type: ignore

1519 -

1520 - image_patch_size = processor.patch_size

1521 - assert isinstance(image_patch_size, int)

1522 -

1523 - return image_patch_size

1524 -

1525 - @cached_property

1526 - def overlap_margins(self) -> tuple[int, int]:

1527 - image_processor = self.processor.image_processor # type: ignore

1528 -

1529 - left_margin, right_margin = image_processor.overlap_margins

1530 - assert isinstance(left_margin, int)

1531 - assert isinstance(right_margin, int)

1532 -

1533 - return left_margin, right_margin

1534 -

1535 - @cached_property

1536 - def bos_token(self) -> str:

1537 - return self.processor.tokenizer.bos_token or self.processor.tokenizer.eos_token

1538 -

1539 - @cached_property

1540 - def image_patch_id(self) -> int:

1541 - return self.hf_config.image_patch_id

1542 -

1543 - @cached_property

1544 - def im_col_id(self) -> int:

1545 - return self.hf_config.image_col_id

1546 -

1547 - @cached_property

1548 - def im_start_id(self) -> int:

1549 - return self.hf_config.image_start_token_id

1550 -

1551 - @cached_property

1552 - def im_end_id(self) -> int:

1553 - return self.hf_config.image_end_token_id

1554 -

1555 - @cached_property

1556 - def low_res_im_start_id(self) -> int:

1557 - return self.hf_config.low_res_image_start_token_id

1558 -

1559 - @cached_property

1560 - def frame_start_id(self) -> int:

1561 - return self.hf_config.frame_start_token_id

1562 -

1563 - @cached_property

1564 - def frame_end_id(self) -> int:

1565 - return self.hf_config.frame_end_token_id

1566 -

1567 - @cached_property

1568 - def im_low_res_id(self) -> int:

1569 - return self.hf_config.image_low_res_id

1570 -

1571 - @cached_property

1572 - def image_placeholder_id(self) -> int:

1573 - return self.vocab[IMAGE_PROMPT]

1574 -

1575 - @cached_property

1576 - def video_placeholder_id(self) -> int:

1577 - return self.vocab[VIDEO_PROMPT]

1578 -

1579 - @cached_property

1580 - def image_token_ids(self) -> list[int]:

1581 - return [

1582 - self.image_patch_id,

1583 - self.im_col_id,

1584 - self.im_start_id,

1585 - self.low_res_im_start_id,

1586 - self.frame_start_id,

1587 - self.im_end_id,

1588 - self.frame_end_id,

1589 - self.im_low_res_id,

1590 - ]

1591 -

1592 - def select_tiling(

1593 - self,

1594 - *,

1595 - image_height: int,

1596 - image_width: int,

1597 - ) -> tuple[int, int]:

1598 - max_crops = self.max_crops

1599 - left_margin, right_margin = self.overlap_margins

1600 - base_image_input_size = self.base_image_input_size

1601 - base_image_input_d = self.image_patch_size

1602 -

1603 - total_margin_pixels = base_image_input_d * (right_margin + left_margin)

1604 - crop_patches = base_image_input_size[0] // base_image_input_d

1605 - crop_window_patches = crop_patches - (right_margin + left_margin)

1606 - crop_window_size = crop_window_patches * base_image_input_d

1607 - tiling_h, tiling_w = select_tiling(

1608 - height=image_height - total_margin_pixels,

1609 - width=image_width - total_margin_pixels,

1610 - patch_size=crop_window_size,

1611 - max_num_patches=max_crops,

1612 - )

1613 -

1614 - return tiling_h, tiling_w

1615 -

1616 - def get_base_grid_size(self, is_video: bool) -> tuple[int, int]:

1617 - base_image_input_size = self.base_image_input_size

1618 -

1619 - return get_patches_grid_size(

1620 - image_h=base_image_input_size[0],

1621 - image_w=base_image_input_size[1],

1622 - patch_size=self.image_patch_size,

1623 - pool_h=self.video_pooling_h if is_video else self.image_pooling_h,

1624 - pool_w=self.video_pooling_w if is_video else self.image_pooling_w,

1625 - )

1626 -

1627 - def get_patches_grid_size(

1628 - self,

1629 - *,

1630 - image_height: int,

1631 - image_width: int,

1632 - ) -> tuple[int, int]:

1633 - left_margin, right_margin = self.overlap_margins

1634 - base_image_input_size = self.base_image_input_size

1635 - base_image_input_d = self.image_patch_size

1636 -

1637 - total_margin_pixels = base_image_input_d * (right_margin + left_margin)

1638 - crop_patches = base_image_input_size[0] // base_image_input_d

1639 - crop_window_patches = crop_patches - (right_margin + left_margin)

1640 - crop_window_size = crop_window_patches * base_image_input_d

1641 -

1642 - tiling_h, tiling_w = self.select_tiling(

1643 - image_height=image_height,

1644 - image_width=image_width,

1645 - )

1646 -

1647 - h, w = [

1648 - tiling_h * crop_window_size + total_margin_pixels,

1649 - tiling_w * crop_window_size + total_margin_pixels,

1650 - ]

1651 - nrows, ncols = get_patches_grid_size(

1652 - image_h=h,

1653 - image_w=w,

1654 - patch_size=base_image_input_d,

1655 - pool_h=self.image_pooling_h,

1656 - pool_w=self.image_pooling_w,

1657 - )

1658 -

1659 - return nrows, ncols

1660 -

1661 - def __call__(

1662 - self,

1663 - text: TextInput | list[TextInput] | None = None,

1664 - images: ImageInput | None = None,

1665 - videos: VideoInput | None = None,

1666 - return_tensors: str | TensorType = None,

1667 - **kwargs: object,

1668 - ) -> BatchFeature:

1669 - inputs = [text]

1670 - images = exif_transpose(images)

1671 - if getattr(self.processor, "image_processor", None) is not None:

1672 - inputs.append(images)

1673 - if getattr(self.processor, "video_processor", None) is not None:

1674 - inputs.append(videos)

1675 - outputs = self.processor( # type: ignore

1676 - *inputs,

1677 - return_tensors=return_tensors,

1678 - **kwargs,

1679 - )

1680 -

1681 - # revert insert bos token

1682 - if outputs["input_ids"][0, 0] == self.vocab[self.bos_token]:

1683 - outputs["input_ids"] = outputs["input_ids"][:, 1:]

1684 -

1685 - if images is None:

1686 - images = []

1687 - if not isinstance(images, list):

1688 - images = [images]

1689 -

1690 - if videos is None:

1691 - videos = []

1692 - if not isinstance(videos, list):

1693 - videos = [videos]

1694 -

1695 - assert len(videos) in {0, 1}, "At most one video is supported for Molmo2"

1696 -

1697 - _attention_mask: torch.Tensor = outputs.pop("attention_mask")

1698 - _token_type_ids: torch.Tensor = outputs.pop("token_type_ids", None)

1699 -

1700 - if len(images) > 0:

1701 - # For each image: tiling_h * tiling_w + global view

1702 - num_crops = []

1703 - for image in images:

1704 - image_size = get_image_size(image)

1705 - tiling = self.select_tiling(

1706 - image_height=image_size.height,

1707 - image_width=image_size.width,

1708 - )

1709 - num_crops.append(np.prod(tiling) + 1)

1710 -

1711 - assert sum(num_crops) == len(outputs["pixel_values"])

1712 - assert sum(num_crops) == outputs["image_num_crops"].sum().item()

1713 - image_grids: torch.Tensor = outputs.pop("image_grids")

1714 - image_num_pooled_patches: torch.Tensor = image_grids[:, :2].prod(

1715 - dim=1

1716 - ) + image_grids[:, 2:].prod(dim=1)

1717 - outputs["image_num_pooled_patches"] = image_num_pooled_patches

1718 - n_patches = outputs["pixel_values"].shape[1]

1719 - outputs["image_num_patches"] = outputs["image_num_crops"] * n_patches

1720 - image_tokens, num_image_tokens = build_flat_image_bool_length(

1721 - image_grids,

1722 - self.image_patch_id,

1723 - self.low_res_im_start_id,

1724 - self.im_start_id,

1725 - self.im_col_id,

1726 - self.im_end_id,

1727 - )

1728 - outputs["image_tokens"] = image_tokens

1729 - outputs["num_image_tokens"] = num_image_tokens

1730 -

1731 - if len(videos) > 0:

1732 - video_grids: torch.Tensor = outputs.pop("video_grids")

1733 - assert video_grids[:, 0].sum() == len(outputs["pixel_values_videos"])

1734 - outputs["video_num_crops"] = video_grids[:, 0]

1735 - outputs["video_num_pooled_patches"] = video_grids.prod(dim=1)

1736 - n_patches = outputs["pixel_values_videos"].shape[1]

1737 - outputs["video_num_patches"] = outputs["video_num_crops"] * n_patches

1738 - video_tokens, num_video_tokens = build_flat_video_bool_length(

1739 - video_grids,

1740 - self.image_patch_id,

1741 - self.frame_start_id,

1742 - self.frame_end_id,

1743 - )

1744 - outputs["video_tokens"] = video_tokens

1745 - outputs["num_video_tokens"] = num_video_tokens

1746 -

1747 - return BatchFeature(outputs)

1748 -

1749 -

1446 def get_candidate_target_fps(

1447 video_fps: int | float,

1448 sampling_fps: int | float,

@@ -1856,36 +1552,101 @@ class Molmo2ProcessingInfo(BaseProcessingInfo):

1552 expected_hidden_size=self._get_expected_hidden_size(),

1553 )

1554

1859 - def get_hf_processor(self, **kwargs: object) -> Molmo2ProcessorWrapper:

1860 - processor = self.ctx.get_hf_processor(**kwargs)

1861 - hf_config = self.ctx.get_hf_config()

1862 - return Molmo2ProcessorWrapper(processor, hf_config)

1863 -

1555 def get_supported_mm_limits(self) -> Mapping[str, int | None]:

1556 return {"image": None, "video": 1}

1557

1558 + def select_tiling(

1559 + self,

1560 + *,

1561 + image_width: int,

1562 + image_height: int,

1563 + image_processor: BaseImageProcessor,

1564 + ) -> tuple[int, int]:

1565 + max_crops = image_processor.max_crops

1566 + left_margin, right_margin = image_processor.overlap_margins

1567 + base_image_input_d = image_processor.patch_size

1568 +

1569 + total_margin_pixels = base_image_input_d * (right_margin + left_margin)

1570 + crop_patches = image_processor.size["height"] // base_image_input_d

1571 + crop_window_patches = crop_patches - (right_margin + left_margin)

1572 + crop_window_size = crop_window_patches * base_image_input_d

1573 + tiling_h, tiling_w = select_tiling(

1574 + height=image_height - total_margin_pixels,

1575 + width=image_width - total_margin_pixels,

1576 + patch_size=crop_window_size,

1577 + max_num_patches=max_crops,

1578 + )

1579 +

1580 + return tiling_w, tiling_h

1581 +

1582 + def get_base_grid_size(

1583 + self,

1584 + image_processor: BaseImageProcessor | BaseVideoProcessor,

1585 + ) -> tuple[int, int]:

1586 + nrows, ncols = get_patches_grid_size(

1587 + image_h=image_processor.size["height"],

1588 + image_w=image_processor.size["width"],

1589 + patch_size=image_processor.patch_size,

1590 + pool_h=image_processor.pooling_size[0],

1591 + pool_w=image_processor.pooling_size[1],

1592 + )

1593 +

1594 + return ncols, nrows

1595 +

1596 + def get_patches_grid_size(

1597 + self,

1598 + *,

1599 + image_width: int,

1600 + image_height: int,

1601 + image_processor: BaseImageProcessor,

1602 + ) -> tuple[int, int]:

1603 + left_margin, right_margin = image_processor.overlap_margins

1604 + base_image_input_d = image_processor.patch_size

1605 +

1606 + total_margin_pixels = base_image_input_d * (right_margin + left_margin)

1607 + crop_patches = image_processor.size["height"] // base_image_input_d

1608 + crop_window_patches = crop_patches - (right_margin + left_margin)

1609 + crop_window_size = crop_window_patches * base_image_input_d

1610 +

1611 + tiling_w, tiling_h = self.select_tiling(

1612 + image_height=image_height,

1613 + image_width=image_width,

1614 + image_processor=image_processor,

1615 + )

1616 +

1617 + nrows, ncols = get_patches_grid_size(

1618 + image_h=tiling_h * crop_window_size + total_margin_pixels,

1619 + image_w=tiling_w * crop_window_size + total_margin_pixels,

1620 + patch_size=base_image_input_d,

1621 + pool_h=image_processor.pooling_size[0],

1622 + pool_w=image_processor.pooling_size[1],

1623 + )

1624 +

1625 + return ncols, nrows

1626 +

1627 def get_num_image_tokens(

1628 self,

1629 *,

1630 image_height: int,

1631 image_width: int,

1872 - processor: Molmo2ProcessorWrapper,

1632 + processor: ProcessorMixin,

1633 ) -> int:

1874 - hf_processor = processor.processor

1634 + image_processor = processor.image_processor

1635

1876 - resize_nrows, resize_cols = processor.get_base_grid_size(is_video=False)

1636 + resize_ncols, resize_nrows = self.get_base_grid_size(image_processor)

1637 # start/end tokens + image patch token + col tokens

1878 - if hf_processor.use_single_crop_col_tokens is not None:

1879 - use_col_tokens = hf_processor.use_single_crop_col_tokens

1638 + if processor.use_single_crop_col_tokens is not None:

1639 + use_col_tokens = processor.use_single_crop_col_tokens

1640 else:

1881 - use_col_tokens = hf_processor.image_use_col_tokens

1882 - extra = 2 + resize_nrows * (resize_cols + int(use_col_tokens))

1883 - overlap_nrows, overlap_ncols = processor.get_patches_grid_size(

1641 + use_col_tokens = processor.image_use_col_tokens

1642 + extra = 2 + resize_nrows * (resize_ncols + int(use_col_tokens))

1643 + overlap_ncols, overlap_nrows = self.get_patches_grid_size(

1644 image_height=image_height,

1645 image_width=image_width,

1646 + image_processor=image_processor,

1647 )

1648 joint = 2 + overlap_nrows * (

1888 - overlap_ncols + int(hf_processor.image_use_col_tokens)

1649 + overlap_ncols + int(processor.image_use_col_tokens)

1650 )

1651

1652 return extra + joint

@@ -1894,28 +1655,28 @@ class Molmo2ProcessingInfo(BaseProcessingInfo):

1655 self,

1656 *,

1657 num_frames: int,

1897 - processor: Molmo2ProcessorWrapper,

1658 + processor: ProcessorMixin,

1659 ) -> int:

1899 - resize_nrows, resize_cols = processor.get_base_grid_size(is_video=True)

1660 + video_processor = processor.video_processor

1661 +

1662 + resize_ncols, resize_nrows = self.get_base_grid_size(video_processor)

1663 # start/end tokens

1901 - extra = 2 + resize_nrows * (

1902 - resize_cols + int(processor.processor.video_use_col_tokens)

1903 - )

1664 + extra = 2 + resize_nrows * (resize_ncols + int(processor.video_use_col_tokens))

1665 return num_frames * extra

1666

1667 def get_image_size_with_most_features(self) -> ImageSize:

1668 processor = self.get_hf_processor()

1669 + image_processor = processor.image_processor

1670

1909 - left_margin, right_margin = processor.overlap_margins

1910 - base_image_input_size = processor.base_image_input_size

1911 - base_image_input_d = processor.image_patch_size

1671 + left_margin, right_margin = image_processor.overlap_margins

1672 + base_image_input_d = image_processor.patch_size

1673

1674 total_margin_pixels = base_image_input_d * (right_margin + left_margin)

1914 - crop_patches = base_image_input_size[0] // base_image_input_d

1675 + crop_patches = image_processor.size["height"] // base_image_input_d

1676 crop_window_patches = crop_patches - (right_margin + left_margin)

1677 crop_window_size = crop_window_patches * base_image_input_d

1678

1918 - tilings = get_candidate_tilings(processor.max_crops)

1679 + tilings = get_candidate_tilings(image_processor.max_crops)

1680 largest_feature_size, largest_feature_pinpoint = 0, None

1681

1682 for hr, wr in tilings:

@@ -1939,7 +1700,7 @@ class Molmo2ProcessingInfo(BaseProcessingInfo):

1700 def _get_max_video_frames(

1701 self,

1702 max_tokens: int,

1942 - processor: Molmo2ProcessorWrapper,

1703 + processor: ProcessorMixin,

1704 ) -> int:

1705 num_tokens_per_frame = self.get_num_video_tokens(

1706 num_frames=1,

@@ -1954,7 +1715,8 @@ class Molmo2ProcessingInfo(BaseProcessingInfo):

1715 mm_counts: Mapping[str, int],

1716 ) -> int:

1717 processor = self.get_hf_processor()

1957 - video_processor = processor.processor.video_processor

1718 + video_processor = processor.video_processor

1719 +

1720 num_frames = video_processor.num_frames

1721 max_videos = mm_counts.get("video", 0)

1722 max_total_frames = self._get_max_video_frames(seq_len, processor)

@@ -2030,7 +1792,9 @@ class Molmo2ProcessingInfo(BaseProcessingInfo):

1792 metadata: dict[str, Any],

1793 do_sample_frames: bool | None = None,

1794 ) -> list[float]:

2033 - video_processor = self.get_hf_processor().processor.video_processor

1795 + processor = self.get_hf_processor()

1796 + video_processor = processor.video_processor

1797 +

1798 # metadata["fps"] refers to the true fps of the input video.

1799 video_fps = metadata["fps"]

1800 frames_indices = metadata.get("frames_indices")

@@ -2104,7 +1868,7 @@ class Molmo2DummyInputsBuilder(BaseDummyInputsBuilder[Molmo2ProcessingInfo]):

1868

1869 if num_videos > 0:

1870 processor = self.info.get_hf_processor()

2107 - base_image_input_size = processor.base_image_input_size

1871 + video_size = processor.video_processor.size

1872 target_num_frames = self.info.get_num_frames_with_most_features(

1873 seq_len, mm_counts

1874 )

@@ -2131,8 +1895,8 @@ class Molmo2DummyInputsBuilder(BaseDummyInputsBuilder[Molmo2ProcessingInfo]):

1895 target_num_frames = min(target_num_frames, num_frames_override)

1896

1897 dummy_videos = self._get_dummy_videos(

2134 - width=base_image_input_size[1],

2135 - height=base_image_input_size[0],

1898 + width=video_size["width"],

1899 + height=video_size["height"],

1900 num_frames=target_num_frames,

1901 num_videos=num_videos,

1902 )

@@ -2174,10 +1938,10 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

1938 prompt_tokens: list[int],

1939 ) -> list[int]:

1940 processor = self.info.get_hf_processor()

2177 - tokenizer = processor.processor.tokenizer

1941 + tokenizer = processor.tokenizer

1942 bos_token_id = tokenizer.bos_token_id or tokenizer.eos_token_id

1943

2180 - if len(prompt_tokens) > 0 and prompt_tokens[0] != bos_token_id:

1944 + if len(prompt_tokens) == 0 or prompt_tokens[0] != bos_token_id:

1945 # Prepend the bos token to the prompt tokens

1946 prompt_tokens = [bos_token_id] + prompt_tokens

1947

@@ -2191,9 +1955,26 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

1955 tok_kwargs: Mapping[str, object],

1956 ) -> BatchFeature:

1957 mm_data = dict(mm_data)

2194 - processor = self.info.get_hf_processor(**mm_kwargs)

1958 +

1959 + hf_config = self.info.get_hf_config()

1960 + hf_processor = self.info.get_hf_processor(**mm_kwargs)

1961 +

1962 + def patched_call(text=None, images=None, videos=None, **kwargs) -> BatchFeature:

1963 + res = hf_processor(text=text, images=images, videos=videos, **kwargs)

1964 +

1965 + # Molmo2Processor.insert_bos results in float outputs

1966 + # if the input text is empty

1967 + if not text:

1968 + res["input_ids"] = res["input_ids"].long()

1969 +

1970 + return res

1971 +

1972 + tokenizer = hf_processor.tokenizer

1973 + image_processor = hf_processor.image_processor

1974

1975 if videos := mm_data.pop("videos", []):

1976 + bos_token_id = tokenizer.bos_token_id or tokenizer.eos_token_id

1977 +

1978 pixel_values_videos_lst = []

1979 video_token_pooling_lst = []

1980 video_num_crops_lst = []

@@ -2228,18 +2009,32 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

2009 video_mm_data["videos"] = [[video_array]]

2010 video_mm_data["video_metadata"] = [[metadata]]

2011

2231 - video_outputs = super()._call_hf_processor(

2232 - prompt=VIDEO_PROMPT,

2233 - mm_data=video_mm_data,

2234 - mm_kwargs=video_mm_kwargs,

2235 - tok_kwargs=tok_kwargs,

2012 + video_outputs = self.info.ctx.call_hf_processor(

2013 + patched_call,

2014 + dict(text=VIDEO_PROMPT, **video_mm_data),

2015 + dict(**video_mm_kwargs, **tok_kwargs),

2016 )

2017 +

2018 input_ids = video_outputs.pop("input_ids")

2238 - video_string = processor.processor.tokenizer.batch_decode(input_ids)[0]

2239 - prompt = prompt.replace(

2240 - VIDEO_PROMPT,

2241 - video_string,

2242 - 1,

2019 + if input_ids[0, 0] == bos_token_id:

2020 + input_ids = input_ids[:, 1:]

2021 +

2022 + video_string = tokenizer.batch_decode(input_ids)[0]

2023 + prompt = prompt.replace(VIDEO_PROMPT, video_string, 1)

2024 +

2025 + video_grids = video_outputs.pop("video_grids")

2026 + assert video_grids[:, 0].sum() == len(

2027 + video_outputs["pixel_values_videos"]

2028 + )

2029 +

2030 + video_outputs["video_num_crops"] = video_grids[:, 0]

2031 + video_outputs["video_num_pooled_patches"] = video_grids.prod(dim=1)

2032 + n_patches = video_outputs["pixel_values_videos"].shape[1]

2033 + video_outputs["video_num_patches"] = (

2034 + video_outputs["video_num_crops"] * n_patches

2035 + )

2036 + (video_outputs["video_tokens"], video_outputs["num_video_tokens"]) = (

2037 + build_flat_video_bool_length(video_grids, hf_config)

2038 )

2039

2040 pixel_values_videos_lst.append(video_outputs["pixel_values_videos"])

@@ -2252,7 +2047,7 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

2047 video_tokens_lst.append(video_outputs["video_tokens"])

2048 num_video_tokens_lst.append(video_outputs["num_video_tokens"])

2049

2255 - video_outputs = dict(

2050 + all_video_outputs = dict(

2051 pixel_values_videos=torch.cat(pixel_values_videos_lst),

2052 video_token_pooling=torch.cat(video_token_pooling_lst),

2053 video_num_crops=torch.cat(video_num_crops_lst),

@@ -2262,30 +2057,50 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

2057 num_video_tokens=torch.cat(num_video_tokens_lst),

2058 )

2059 else:

2265 - video_outputs = dict()

2060 + all_video_outputs = dict()

2061

2267 - processed_outputs = super()._call_hf_processor(

2268 - prompt=prompt,

2269 - mm_data=mm_data,

2270 - mm_kwargs=mm_kwargs,

2271 - tok_kwargs=tok_kwargs,

2062 + processed_outputs = self.info.ctx.call_hf_processor(

2063 + patched_call,

2064 + dict(text=prompt, **mm_data),

2065 + dict(**mm_kwargs, **tok_kwargs),

2066 )

2067

2274 - bos_token_id = processor.vocab[processor.bos_token]

2275 - input_ids = processed_outputs["input_ids"]

2276 - # add bos token back to prompt start

2277 - if input_ids.numel() > 0 and input_ids[0, 0] != bos_token_id:

2278 - bos_token_id_tensor = torch.tensor(

2279 - [[bos_token_id]], device=input_ids.device, dtype=input_ids.dtype

2280 - )

2281 - processed_outputs["input_ids"] = torch.concat(

2282 - [bos_token_id_tensor, input_ids], dim=1

2068 + if (images := mm_data.get("images")) is not None:

2069 + mm_items = self.info.parse_mm_data({"image": images}, validate=False)

2070 + parsed_images = mm_items.get_items("image", ImageProcessorItems)

2071 + image_sizes = [

2072 + parsed_images.get_image_size(i) for i in range(len(parsed_images))

2073 + ]

2074 +

2075 + # For each image: tiling_h * tiling_w + global view

2076 + tilings = [

2077 + self.info.select_tiling(

2078 + image_width=image_size.width,

2079 + image_height=image_size.height,

2080 + image_processor=image_processor,

2081 + )

2082 + for image_size in image_sizes

2083 + ]

2084 + num_crops = torch.tensor(tilings).prod(-1) + 1

2085 + assert sum(num_crops) == len(processed_outputs["pixel_values"])

2086 + assert sum(num_crops) == processed_outputs["image_num_crops"].sum().item()

2087 +

2088 + image_grids = processed_outputs.pop("image_grids")

2089 + image_num_pooled_patches = image_grids[:, :2].prod(dim=1) + image_grids[

2090 + :, 2:

2091 + ].prod(dim=1)

2092 +

2093 + processed_outputs["image_num_pooled_patches"] = image_num_pooled_patches

2094 + n_patches = processed_outputs["pixel_values"].shape[1]

2095 + processed_outputs["image_num_patches"] = (

2096 + processed_outputs["image_num_crops"] * n_patches

2097 )

2284 - combined_outputs = dict(

2285 - processed_outputs,

2286 - **video_outputs,

2287 - )

2288 - return BatchFeature(combined_outputs)

2098 + (

2099 + processed_outputs["image_tokens"],

2100 + processed_outputs["num_image_tokens"],

2101 + ) = build_flat_image_bool_length(image_grids, hf_config)

2102 +

2103 + return BatchFeature({**processed_outputs, **all_video_outputs})

2104

2105 def _get_mm_fields_config(

2106 self,

@@ -2338,41 +2153,65 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

2153 hf_processor_mm_kwargs: Mapping[str, object],

2154 out_mm_kwargs: MultiModalKwargsItems,

2155 ) -> Sequence[PromptUpdate]:

2156 + hf_config = self.info.get_hf_config()

2157 + img_patch_id = hf_config.image_patch_id

2158 + img_col_id = hf_config.image_col_id

2159 + img_start_id = hf_config.image_start_token_id

2160 + img_end_id = hf_config.image_end_token_id

2161 + low_res_im_start_id = hf_config.low_res_image_start_token_id

2162 + frame_start_id = hf_config.frame_start_token_id

2163 + frame_end_id = hf_config.frame_end_token_id

2164 + im_low_res_id = hf_config.image_low_res_id

2165 +

2166 + emb_tok_ids = [

2167 + img_patch_id,

2168 + img_col_id,

2169 + img_start_id,

2170 + low_res_im_start_id,

2171 + frame_start_id,

2172 + img_end_id,

2173 + frame_end_id,

2174 + im_low_res_id,

2175 + ]

2176 +

2177 processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

2342 - img_patch_id = processor.image_patch_id

2343 - img_col_id = processor.im_col_id

2344 - img_start_id = processor.im_start_id

2345 - img_end_id = processor.im_end_id

2346 - image_use_col_tokens = processor.processor.image_use_col_tokens

2347 - use_single_crop_col_tokens = processor.processor.use_single_crop_col_tokens

2348 - use_single_crop_start_token = processor.processor.use_single_crop_start_token

2349 - video_use_col_tokens = processor.processor.video_use_col_tokens

2350 - use_frame_special_tokens = processor.processor.use_frame_special_tokens

2351 -

2352 - def get_image_replacement_molmo2(item_idx: int) -> list[int]:

2178 + image_use_col_tokens = processor.image_use_col_tokens

2179 + use_single_crop_col_tokens = processor.use_single_crop_col_tokens

2180 + use_single_crop_start_token = processor.use_single_crop_start_token

2181 + video_use_col_tokens = processor.video_use_col_tokens

2182 + use_frame_special_tokens = processor.use_frame_special_tokens

2183 +

2184 + tokenizer = processor.tokenizer

2185 + vocab = tokenizer.get_vocab()

2186 +

2187 + image_processor = processor.image_processor

2188 + video_processor = processor.video_processor

2189 +

2190 + def get_image_replacement_molmo2(item_idx: int):

2191 images = mm_items.get_items("image", ImageProcessorItems)

2192 image = images.get(item_idx)

2193 image = exif_transpose(image)

2194

2357 - resize_nrows, resize_cols = processor.get_base_grid_size(is_video=False)

2195 + resize_ncols, resize_nrows = self.info.get_base_grid_size(image_processor)

2196 if use_single_crop_col_tokens is not None:

2197 use_col_tokens = use_single_crop_col_tokens

2198 else:

2199 use_col_tokens = image_use_col_tokens

2200 if use_single_crop_start_token:

2363 - start_id = processor.low_res_im_start_id

2201 + start_id = low_res_im_start_id

2202 else:

2203 start_id = img_start_id

2366 - extra_row = [img_patch_id] * resize_cols + [img_col_id] * int(

2204 + extra_row = [img_patch_id] * resize_ncols + [img_col_id] * int(

2205 use_col_tokens

2206 )

2207 extra_joint = [start_id] + extra_row * resize_nrows + [img_end_id]

2208

2209 image_size = get_image_size(image)

2210

2373 - nrows, ncols = processor.get_patches_grid_size(

2211 + ncols, nrows = self.info.get_patches_grid_size(

2212 image_height=image_size.height,

2213 image_width=image_size.width,

2214 + image_processor=image_processor,

2215 )

2216

2217 joint_row = [img_patch_id] * ncols + [img_col_id] * int(

@@ -2381,21 +2220,18 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

2220 joint = [img_start_id] + joint_row * nrows + [img_end_id]

2221 img_token_ids = extra_joint + joint

2222

2384 - return PromptUpdateDetails.select_token_ids(

2385 - img_token_ids,

2386 - processor.image_token_ids,

2387 - )

2223 + return PromptUpdateDetails.select_token_ids(img_token_ids, emb_tok_ids)

2224

2389 - def get_video_replacement_molmo2(item_idx: int) -> list[int]:

2225 + def get_video_replacement_molmo2(item_idx: int):

2226 video, metadata = mm_items["video"][item_idx]

2227 do_sample_frames = hf_processor_mm_kwargs.get("do_sample_frames")

2228

2229 timestamps = self.info._get_video_second_idx(metadata, do_sample_frames)

2394 - nrows, ncols = processor.get_base_grid_size(is_video=True)

2230 + ncols, nrows = self.info.get_base_grid_size(video_processor)

2231

2232 if use_frame_special_tokens:

2397 - start_id = processor.frame_start_id

2398 - end_id = processor.frame_end_id

2233 + start_id = frame_start_id

2234 + end_id = frame_end_id

2235 else:

2236 start_id = img_start_id

2237 end_id = img_end_id

@@ -2408,7 +2244,7 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

2244 prev_space + f"{frame_time:.1f} "

2245 ) # explicit whitespace before/after image tokens

2246

2411 - img_token_ids += processor.processor.tokenizer.encode(

2247 + img_token_ids += tokenizer.encode(

2248 frame_prefix,

2249 add_special_tokens=False,

2250 )

@@ -2419,10 +2255,7 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

2255 joint = [start_id] + nrows * joint_row + [end_id]

2256 img_token_ids += joint

2257

2422 - return PromptUpdateDetails.select_token_ids(

2423 - img_token_ids,

2424 - processor.image_token_ids,

2425 - )

2258 + return PromptUpdateDetails.select_token_ids(img_token_ids, emb_tok_ids)

2259

2260 return [

2261 PromptReplacement(

@@ -2432,7 +2265,7 @@ class Molmo2MultiModalProcessor(BaseMultiModalProcessor[Molmo2ProcessingInfo]):

2265 )

2266 for modality, target, replacement_fn in zip(

2267 ["image", "video"],

2435 - [processor.image_placeholder_id, processor.video_placeholder_id],

2268 + [vocab[IMAGE_PROMPT], vocab[VIDEO_PROMPT]],

2269 [get_image_replacement_molmo2, get_video_replacement_molmo2],

2270 )

2271 ]

[Refactor] Remove Molmo2 processor wrapper (#36667)