Aluode/PerceptionLabPortable

1#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2#           This file was automatically generated from src/transformers/models/aya_vision/modular_aya_vision.py.3#               Do NOT edit this file manually as any edits will be overwritten by the generation of4#             the file from the modular. If any change should be done, please apply the change to the5#                          modular_aya_vision.py file directly. One of our CI enforces this.6#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# coding=utf-88# Copyright 2025 the Cohere Inc. team. All rights reserved.9#10# Licensed under the Apache License, Version 2.0 (the "License");11# you may not use this file except in compliance with the License.12# You may obtain a copy of the License at13#14#     http://www.apache.org/licenses/LICENSE-2.015#16# Unless required by applicable law or agreed to in writing, software17# distributed under the License is distributed on an "AS IS" BASIS,18# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.19# See the License for the specific language governing permissions and20# limitations under the License.21 22from dataclasses import dataclass23from typing import Optional, Union24 25import torch26from torch import nn27 28from ...activations import ACT2FN29from ...cache_utils import Cache30from ...generation import GenerationMixin31from ...modeling_outputs import BaseModelOutputWithPast, ModelOutput32from ...modeling_utils import PreTrainedModel33from ...processing_utils import Unpack34from ...utils import TransformersKwargs, auto_docstring, can_return_tuple35from ...utils.generic import check_model_inputs36from ..auto import AutoModel37from .configuration_aya_vision import AyaVisionConfig38 39 40class AyaVisionMultiModalProjector(nn.Module):41    def __init__(self, config: AyaVisionConfig):42        super().__init__()43        self.config = config44        self.downsample_factor = config.downsample_factor45        self.alignment_intermediate_size = getattr(46            config, "alignment_intermediate_size", config.text_config.hidden_size47        )48        self.layernorm = nn.LayerNorm(49            config.vision_config.hidden_size * (config.downsample_factor**2), eps=config.adapter_layer_norm_eps50        )51 52        self.linear_1 = nn.Linear(53            config.vision_config.hidden_size * (config.downsample_factor**2),54            self.alignment_intermediate_size,55            bias=True,56        )57 58        self.act = ACT2FN["silu"]  # SwiGLU uses SiLU activation59        # For SwiGLU, project down to half size since we split intermediate dim60        self.linear_2 = nn.Linear(self.alignment_intermediate_size // 2, config.text_config.hidden_size, bias=True)61 62    def forward(self, image_features):63        image_features = self.pixel_shuffle(image_features)64        image_features = self.layernorm(image_features)65        hidden_states = self.linear_1(image_features)66 67        # Split along last dimension and apply SwiGLU68        x, gate = hidden_states.chunk(2, dim=-1)69        hidden_states = self.act(gate) * x70 71        hidden_states = self.linear_2(hidden_states)72        return hidden_states73 74    def pixel_shuffle(self, image_features):  # B, S, D75        batch_size, seq_length, feature_dim = image_features.shape76        height = width = int(seq_length**0.5)77        image_features = image_features.reshape(image_features.shape[0], width, height, -1)78        channels = image_features.shape[-1]79        image_features = image_features.reshape(80            batch_size, width, int(height / self.downsample_factor), int(channels * self.downsample_factor)81        )82        image_features = image_features.permute(0, 2, 1, 3)83        image_features = image_features.reshape(84            batch_size, int(height / self.downsample_factor), int(width / self.downsample_factor), -185        )86        image_features = image_features.permute(0, 2, 1, 3)87        return image_features88 89 90@auto_docstring91class AyaVisionPreTrainedModel(PreTrainedModel):92    config: AyaVisionConfig93    base_model_prefix = ""94    supports_gradient_checkpointing = True95    _skip_keys_device_placement = "past_key_values"96 97    _supports_flash_attn = True98    _supports_sdpa = True99    _can_compile_fullgraph = False100    _supports_flex_attn = True101    _supports_attention_backend = True102    _can_record_outputs = {103        "hidden_states": "DecoderLayer",104        "attentions": "Attention",105    }106 107 108@dataclass109@auto_docstring(110    custom_intro="""111    Base class for AyaVision causal language model (or autoregressive) outputs.112    """113)114class AyaVisionCausalLMOutputWithPast(ModelOutput):115    r"""116    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):117        Language modeling loss (for next-token prediction).118    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):119        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).120    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):121        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).122 123        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see124        `past_key_values` input) to speed up sequential decoding.125    image_hidden_states (`torch.FloatTensor`, *optional*):126        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.127        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.128    """129 130    loss: Optional[torch.FloatTensor] = None131    logits: Optional[torch.FloatTensor] = None132    past_key_values: Optional[Cache] = None133    hidden_states: Optional[tuple[torch.FloatTensor]] = None134    attentions: Optional[tuple[torch.FloatTensor]] = None135    image_hidden_states: Optional[torch.FloatTensor] = None136 137 138@dataclass139@auto_docstring(140    custom_intro="""141    Base class for AyaVision outputs, with hidden states and attentions.142    """143)144class AyaVisionModelOutputWithPast(BaseModelOutputWithPast):145    r"""146    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):147        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).148 149        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see150        `past_key_values` input) to speed up sequential decoding.151    image_hidden_states (`torch.FloatTensor`, *optional*):152        A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.153        image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.154    """155 156    image_hidden_states: Optional[torch.FloatTensor] = None157 158 159@auto_docstring(160    custom_intro="""161    The AyaVision model which consists of a vision backbone and a language model, without a language modeling head.162    """163)164class AyaVisionModel(AyaVisionPreTrainedModel):165    _checkpoint_conversion_mapping = {"language_model.model": "language_model"}166 167    def __init__(self, config: AyaVisionConfig):168        super().__init__(config)169        self.vision_tower = AutoModel.from_config(config.vision_config)170 171        self.multi_modal_projector = AyaVisionMultiModalProjector(config)172        self.language_model = AutoModel.from_config(config.text_config)173        self.post_init()174 175    def get_input_embeddings(self):176        return self.language_model.get_input_embeddings()177 178    def set_input_embeddings(self, value):179        self.language_model.set_input_embeddings(value)180 181    def set_decoder(self, decoder):182        self.language_model = decoder183 184    def get_decoder(self):185        return self.language_model186 187    def get_image_features(188        self,189        pixel_values: torch.FloatTensor,190        vision_feature_layer: Optional[Union[int, list[int]]] = None,191        vision_feature_select_strategy: Optional[str] = None,192        **kwargs,193    ):194        """195        Obtains image last hidden states from the vision tower and apply multimodal projection.196 197        Args:198            pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`):199               The tensors corresponding to the input images.200            vision_feature_layer (`Union[int, list[int]]`, *optional*):201                The index of the layer to select the vision feature. If multiple indices are provided,202                the vision feature of the corresponding indices will be concatenated to form the203                vision features.204            vision_feature_select_strategy (`str`, *optional*):205                The feature selection strategy used to select the vision feature from the vision backbone.206                Can be one of `"default"` or `"full"`207        Returns:208            image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`).209        """210        vision_feature_layer = (211            vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer212        )213        vision_feature_select_strategy = (214            vision_feature_select_strategy215            if vision_feature_select_strategy is not None216            else self.config.vision_feature_select_strategy217        )218 219        if vision_feature_select_strategy not in ["default", "full"]:220            raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}")221 222        kwargs = {k: v for k, v in kwargs.items() if v is not None}223        # this is not memory efficient at all (output_hidden_states=True) will save all the hidden states.224        image_outputs = self.vision_tower(pixel_values, output_hidden_states=True, **kwargs)225 226        # If we have one vision feature layer, return the corresponding hidden states,227        # otherwise, select the hidden states of each feature layer and concatenate them228        if isinstance(vision_feature_layer, int):229            selected_image_feature = image_outputs.hidden_states[vision_feature_layer]230            if vision_feature_select_strategy == "default":231                selected_image_feature = selected_image_feature[:, 1:]232        else:233            hs_pool = [image_outputs.hidden_states[layer_idx] for layer_idx in vision_feature_layer]234            # For default; crop CLS from each hidden state in the hidden state pool235            if vision_feature_select_strategy == "default":236                hs_pool = [hs[:, 1:] for hs in hs_pool]237            selected_image_feature = torch.cat(hs_pool, dim=-1)238 239        image_features = self.multi_modal_projector(selected_image_feature)240        return image_features241 242    def get_placeholder_mask(243        self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor244    ):245        """246        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is247        equal to the length of multimodal features. If the lengths are different, an error is raised.248        """249        if input_ids is None:250            special_image_mask = inputs_embeds == self.get_input_embeddings()(251                torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)252            )253            special_image_mask = special_image_mask.all(-1)254        else:255            special_image_mask = input_ids == self.config.image_token_id256 257        n_image_tokens = special_image_mask.sum()258        special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)259        n_image_features = image_features.shape[0] * image_features.shape[1]260        if inputs_embeds[special_image_mask].numel() != image_features.numel():261            raise ValueError(262                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"263            )264        return special_image_mask265 266    @check_model_inputs()267    @auto_docstring268    def forward(269        self,270        input_ids: Optional[torch.LongTensor] = None,271        pixel_values: Optional[torch.FloatTensor] = None,272        attention_mask: Optional[torch.Tensor] = None,273        position_ids: Optional[torch.LongTensor] = None,274        past_key_values: Optional[Cache] = None,275        inputs_embeds: Optional[torch.FloatTensor] = None,276        vision_feature_layer: Optional[Union[int, list[int]]] = None,277        vision_feature_select_strategy: Optional[str] = None,278        use_cache: Optional[bool] = None,279        cache_position: Optional[torch.LongTensor] = None,280        **kwargs: Unpack[TransformersKwargs],281    ) -> Union[tuple, AyaVisionModelOutputWithPast]:282        vision_feature_layer = (283            vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer284        )285        vision_feature_select_strategy = (286            vision_feature_select_strategy287            if vision_feature_select_strategy is not None288            else self.config.vision_feature_select_strategy289        )290 291        if (input_ids is None) ^ (inputs_embeds is not None):292            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")293 294        if inputs_embeds is None:295            inputs_embeds = self.get_input_embeddings()(input_ids)296 297        if pixel_values is not None:298            image_features = self.get_image_features(299                pixel_values=pixel_values,300                vision_feature_layer=vision_feature_layer,301                vision_feature_select_strategy=vision_feature_select_strategy,302            )303            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)304            special_image_mask = self.get_placeholder_mask(305                input_ids, inputs_embeds=inputs_embeds, image_features=image_features306            )307            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)308 309        outputs = self.language_model(310            attention_mask=attention_mask,311            position_ids=position_ids,312            past_key_values=past_key_values,313            inputs_embeds=inputs_embeds,314            use_cache=use_cache,315            cache_position=cache_position,316            **kwargs,317        )318 319        return AyaVisionModelOutputWithPast(320            last_hidden_state=outputs.last_hidden_state,321            past_key_values=outputs.past_key_values,322            hidden_states=outputs.hidden_states,323            attentions=outputs.attentions,324            image_hidden_states=image_features if pixel_values is not None else None,325        )326 327 328@auto_docstring(329    custom_intro="""330    The AYA_VISION model which consists of a vision backbone and a language model.331    """332)333class AyaVisionForConditionalGeneration(AyaVisionPreTrainedModel, GenerationMixin):334    _checkpoint_conversion_mapping = {335        "^language_model.model": "model.language_model",336        "^vision_tower": "model.vision_tower",337        "^multi_modal_projector": "model.multi_modal_projector",338        "^language_model.lm_head": "lm_head",339    }340    _tied_weights_keys = ["lm_head.weight"]341 342    def __init__(self, config: AyaVisionConfig):343        super().__init__(config)344        self.model = AyaVisionModel(config)345        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)346        self.post_init()347 348    def get_input_embeddings(self):349        return self.model.get_input_embeddings()350 351    def set_input_embeddings(self, value):352        self.model.set_input_embeddings(value)353 354    def get_output_embeddings(self) -> nn.Module:355        return self.lm_head356 357    def set_decoder(self, decoder):358        self.model.set_decoder(decoder)359 360    def get_decoder(self):361        return self.model.get_decoder()362 363    def get_image_features(364        self,365        pixel_values: torch.FloatTensor,366        vision_feature_layer: Optional[Union[int, list[int]]] = None,367        vision_feature_select_strategy: Optional[str] = None,368        **kwargs,369    ):370        return self.model.get_image_features(371            pixel_values=pixel_values,372            vision_feature_layer=vision_feature_layer,373            vision_feature_select_strategy=vision_feature_select_strategy,374            **kwargs,375        )376 377    # Make modules available through conditional class for BC378    @property379    def language_model(self):380        return self.model.language_model381 382    @property383    def vision_tower(self):384        return self.model.vision_tower385 386    @property387    def multi_modal_projector(self):388        return self.model.multi_modal_projector389 390    @can_return_tuple391    @auto_docstring392    def forward(393        self,394        input_ids: Optional[torch.LongTensor] = None,395        pixel_values: Optional[torch.FloatTensor] = None,396        attention_mask: Optional[torch.Tensor] = None,397        position_ids: Optional[torch.LongTensor] = None,398        past_key_values: Optional[Cache] = None,399        inputs_embeds: Optional[torch.FloatTensor] = None,400        vision_feature_layer: Optional[Union[int, list[int]]] = None,401        vision_feature_select_strategy: Optional[str] = None,402        labels: Optional[torch.LongTensor] = None,403        cache_position: Optional[torch.LongTensor] = None,404        logits_to_keep: Union[int, torch.Tensor] = 0,405        image_sizes: Optional[torch.Tensor] = None,406        **kwargs: Unpack[TransformersKwargs],407    ) -> Union[tuple, AyaVisionCausalLMOutputWithPast]:408        r"""409        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):410            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,411            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored412            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.413 414        Example:415 416        ```python417        >>> from transformers import AutoProcessor, AyaVisionForConditionalGeneration418        >>> import torch419 420        >>> torch_device = "cuda:0"421        >>> processor = AutoProcessor.from_pretrained("CohereForAI/aya-vision-8b", use_fast=True)422        >>> model = AyaVisionForConditionalGeneration.from_pretrained("CohereForAI/aya-vision-8b", device_map=torch_device)423 424        >>> messages = [425        ...     {426        ...         "role": "user",427        ...         "content": [428        ...             {429        ...                 "type": "image",430        ...                 "url": "https://pbs.twimg.com/media/Fx7YvfQWYAIp6rZ?format=jpg&name=medium",431        ...             },432        ...             {"type": "text", "text": "चित्र में लिखा पाठ क्या कहता है?"},433        ...         ],434        ...     }435        ... ]436 437        >>> inputs = processor.apply_chat_template(438        ...     messages, padding=True, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt", device=torch_device439        ... ).to(model.device)440 441        >>> gen_tokens = model.generate(**inputs, max_new_tokens=300, do_sample=True, temperature=0.3)442        >>> processor.tokenizer.decode(gen_tokens[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)443        ```"""444        vision_feature_layer = (445            vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer446        )447        vision_feature_select_strategy = (448            vision_feature_select_strategy449            if vision_feature_select_strategy is not None450            else self.config.vision_feature_select_strategy451        )452 453        outputs = self.model(454            input_ids=input_ids,455            pixel_values=pixel_values,456            attention_mask=attention_mask,457            position_ids=position_ids,458            past_key_values=past_key_values,459            inputs_embeds=inputs_embeds,460            vision_feature_layer=vision_feature_layer,461            vision_feature_select_strategy=vision_feature_select_strategy,462            cache_position=cache_position,463            image_sizes=image_sizes,464            **kwargs,465        )466 467        hidden_states = outputs[0]468        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss469        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep470        logits = self.lm_head(hidden_states[:, slice_indices, :])471 472        loss = None473        if labels is not None:474            loss = self.loss_function(475                logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs476            )477 478        return AyaVisionCausalLMOutputWithPast(479            loss=loss,480            logits=logits,481            past_key_values=outputs.past_key_values,482            hidden_states=outputs.hidden_states,483            attentions=outputs.attentions,484            image_hidden_states=outputs.image_hidden_states,485        )486 487    def prepare_inputs_for_generation(488        self,489        input_ids,490        past_key_values=None,491        inputs_embeds=None,492        pixel_values=None,493        attention_mask=None,494        cache_position=None,495        logits_to_keep=None,496        **kwargs,497    ):498        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model499 500        model_inputs = super().prepare_inputs_for_generation(501            input_ids,502            past_key_values=past_key_values,503            inputs_embeds=inputs_embeds,504            attention_mask=attention_mask,505            cache_position=cache_position,506            logits_to_keep=logits_to_keep,507            **kwargs,508        )509 510        if cache_position[0] == 0:511            # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore512            # Otherwise we need pixel values to be passed to model513            model_inputs["pixel_values"] = pixel_values514 515        return model_inputs516 517 518__all__ = ["AyaVisionForConditionalGeneration", "AyaVisionPreTrainedModel", "AyaVisionModel"]519