PerceptronAI/Isaac-0.5

1"""Portable Isaac-0.5 processor built on checkpoint-native Qwen3-VL assets."""2 3from __future__ import annotations4 5import math6from typing import Any7 8import torch9from transformers import BatchFeature, Qwen3VLProcessor10 11from .tensor_stream import ALL_TYPES, Event, TensorStream, TextType, VectorType, VisionType, create_stream12 13 14class Isaac05Processor(Qwen3VLProcessor):15    """Produce native Qwen inputs plus PR #5 TensorStream inputs for Isaac-0.5."""16 17    def __call__(18        self,19        images: Any | None = None,20        text: str | list[str] | list[list[str]] | None = None,21        videos: Any | None = None,22        **kwargs: Any,23    ) -> BatchFeature:24        vectors = kwargs.pop("vectors", None)25        features = super().__call__(images=images, text=text, videos=videos, **kwargs)26        features["tensor_stream"] = self._build_tensor_stream(features, vectors=vectors)27        return features28 29    def _build_tensor_stream(30        self,31        features: BatchFeature,32        *,33        vectors: torch.Tensor | None,34    ) -> TensorStream:35        input_ids = features["input_ids"]36        if not isinstance(input_ids, torch.Tensor):37            input_ids = torch.as_tensor(input_ids, dtype=torch.long)38        if input_ids.ndim != 2 or input_ids.shape[0] != 1:39            raise ValueError("Isaac05Processor TensorStream output currently requires batch_size=1.")40 41        token_ids = input_ids[0]42        pixel_values = features.get("pixel_values")43        image_grid_thw = features.get("image_grid_thw")44        if pixel_values is not None and not isinstance(pixel_values, torch.Tensor):45            pixel_values = torch.as_tensor(pixel_values)46        if image_grid_thw is not None and not isinstance(image_grid_thw, torch.Tensor):47            image_grid_thw = torch.as_tensor(image_grid_thw, dtype=torch.long)48 49        events: list[Event] = []50        token_start = 051        patch_start = 052        image_index = 053        image_token_id = int(self.tokenizer.convert_tokens_to_ids("<|image_pad|>"))54        merge_size = int(self.image_processor.merge_size)55 56        while token_start < token_ids.numel():57            image_positions = torch.nonzero(token_ids[token_start:] == image_token_id, as_tuple=False)58            if image_positions.numel() == 0:59                self._append_text_event(events, token_ids[token_start:])60                break61 62            image_start = token_start + int(image_positions[0, 0])63            self._append_text_event(events, token_ids[token_start:image_start])64            if pixel_values is None or image_grid_thw is None or image_index >= image_grid_thw.shape[0]:65                raise ValueError("Isaac05Processor image tokens require matching pixel_values and image_grid_thw.")66 67            grid = image_grid_thw[image_index].to(dtype=torch.long)68            temporal, height, width = (int(value) for value in grid.tolist())69            real_patch_count = temporal * height * width70            if height % merge_size or width % merge_size:71                raise ValueError("Isaac05Processor image grid is not divisible by merge_size.")72            virtual_dims = [temporal, height // merge_size, width // merge_size]73            virtual_token_count = math.prod(virtual_dims)74            image_end = image_start + virtual_token_count75            if not torch.all(token_ids[image_start:image_end] == image_token_id):76                raise ValueError("Isaac05Processor image-token run does not match image_grid_thw.")77 78            events.append(79                Event(80                    data=pixel_values[patch_start : patch_start + real_patch_count],81                    time=(float(image_index), float(image_index)),82                    type=VisionType.I,83                    dims_virtual=virtual_dims,84                    dims_real=[temporal, height, width],85                    idx_range=(0, virtual_token_count),86                )87            )88            patch_start += real_patch_count89            image_index += 190            token_start = image_end91 92        if pixel_values is not None and patch_start != pixel_values.shape[0]:93            raise ValueError("Isaac05Processor did not consume every image patch.")94        if image_grid_thw is not None and image_index != image_grid_thw.shape[0]:95            raise ValueError("Isaac05Processor did not consume every image grid.")96 97        if vectors is not None:98            vector_rows = vectors.reshape(-1, vectors.shape[-1]).to(dtype=torch.float32)99            events.append(100                Event(101                    data=vector_rows,102                    time=(float(len(events)), float(len(events))),103                    type=VectorType.vector,104                    dims_virtual=[vector_rows.shape[0]],105                    dims_real=[vector_rows.shape[0]],106                    idx_range=(0, vector_rows.shape[0]),107                )108            )109 110        return TensorStream([create_stream(events, ALL_TYPES, schedule=False)])111 112    @staticmethod113    def _append_text_event(events: list[Event], token_ids: torch.Tensor) -> None:114        if token_ids.numel() == 0:115            return116        events.append(117            Event.from_text_tokens(118                token_ids,119                time=(float(len(events)), float(len(events))),120                type=TextType.text,121            )122        )123