# coding=utf-8 # Copyright 2025 the Cohere Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """PyTorch AyaVision model.""" from typing import List, Optional, Tuple, Union import torch from torch import nn from transformers.models.llava.modeling_llava import ( KwargsForCausalLM, LlavaCausalLMOutputWithPast, LlavaForConditionalGeneration, LlavaModel, LlavaPreTrainedModel, ) from ...activations import ACT2FN from ...processing_utils import Unpack from ...utils import logging from .configuration_aya_vision import AyaVisionConfig logger = logging.get_logger(__name__) class AyaVisionMultiModalProjector(nn.Module): def __init__(self, config: AyaVisionConfig): super().__init__() self.config = config self.downsample_factor = config.downsample_factor self.alignment_intermediate_size = getattr( config, "alignment_intermediate_size", config.text_config.hidden_size ) self.layernorm = nn.LayerNorm( config.vision_config.hidden_size * (config.downsample_factor**2), eps=config.adapter_layer_norm_eps ) self.linear_1 = nn.Linear( config.vision_config.hidden_size * (config.downsample_factor**2), self.alignment_intermediate_size, bias=True, ) self.act = ACT2FN["silu"] # SwiGLU uses SiLU activation # For SwiGLU, project down to half size since we split intermediate dim self.linear_2 = nn.Linear(self.alignment_intermediate_size // 2, config.text_config.hidden_size, bias=True) def forward(self, image_features): image_features = self.pixel_shuffle(image_features) image_features = self.layernorm(image_features) hidden_states = self.linear_1(image_features) # Split along last dimension and apply SwiGLU x, gate = hidden_states.chunk(2, dim=-1) hidden_states = self.act(gate) * x hidden_states = self.linear_2(hidden_states) return hidden_states def pixel_shuffle(self, image_features): # B, S, D batch_size, seq_length, feature_dim = image_features.shape height = width = int(seq_length**0.5) image_features = image_features.reshape(image_features.shape[0], width, height, -1) channels = image_features.shape[-1] image_features = image_features.reshape( batch_size, width, int(height / self.downsample_factor), int(channels * self.downsample_factor) ) image_features = image_features.permute(0, 2, 1, 3) image_features = image_features.reshape( batch_size, int(height / self.downsample_factor), int(width / self.downsample_factor), -1 ) image_features = image_features.permute(0, 2, 1, 3) return image_features class AyaVisionPreTrainedModel(LlavaPreTrainedModel): _supports_quantized_cache = False _supports_static_cache = False def _init_weights(self, module): std = ( self.config.initializer_range if hasattr(self.config, "initializer_range") else self.config.text_config.initializer_range ) if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.LayerNorm): module.weight.data.fill_(1.0) module.bias.data.zero_() class AyaVisionCausalLMOutputWithPast(LlavaCausalLMOutputWithPast): pass class AyaVisionModel(LlavaModel): pass class AyaVisionForConditionalGeneration(LlavaForConditionalGeneration): def forward( self, input_ids: Optional[torch.LongTensor] = None, pixel_values: Optional[torch.FloatTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, vision_feature_layer: Optional[Union[int, List[int]]] = None, vision_feature_select_strategy: Optional[str] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, image_sizes: Optional[torch.Tensor] = None, **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, AyaVisionCausalLMOutputWithPast]: r""" labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Example: ```python >>> from transformers import AutoProcessor, AyaVisionForConditionalGeneration >>> import torch >>> torch_device = "cuda:0" >>> processor = AutoProcessor.from_pretrained("CohereForAI/aya-vision-8b", use_fast=True) >>> model = AyaVisionForConditionalGeneration.from_pretrained("CohereForAI/aya-vision-8b", device_map=torch_device) >>> messages = [ ... { ... "role": "user", ... "content": [ ... { ... "type": "image", ... "url": "https://pbs.twimg.com/media/Fx7YvfQWYAIp6rZ?format=jpg&name=medium", ... }, ... {"type": "text", "text": "चित्र में लिखा पाठ क्या कहता है?"}, ... ], ... } ... ] >>> inputs = processor.apply_chat_template( ... messages, padding=True, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt", device=torch_device ... ).to(model.device) >>> gen_tokens = model.generate(**inputs, max_new_tokens=300, do_sample=True, temperature=0.3) >>> processor.tokenizer.decode(gen_tokens[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) ```""" super().forward( input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, vision_feature_layer=vision_feature_layer, vision_feature_select_strategy=vision_feature_select_strategy, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position, logits_to_keep=logits_to_keep, image_sizes=image_sizes, **kwargs, ) __all__ = ["AyaVisionForConditionalGeneration", "AyaVisionPreTrainedModel", "AyaVisionModel"]