from typing import List, Optional, Tuple, Union import torch from torch import nn from torch.nn import CrossEntropyLoss from transformers import GenerationConfig, Qwen3Config, Qwen3ForCausalLM from transformers.activations import ACT2FN from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.modeling_utils import PreTrainedModel from .adaptor_base import * # noqa: F401,F403 from .adaptor_generic import * # noqa: F401,F403 from .adaptor_mlp import * # noqa: F401,F403 from .adaptor_registry import * # noqa: F401,F403 from .cls_token import * # noqa: F401,F403 from .configuration_vectorllm import ProjectorConfig, VectorLLMConfig from .common import * # noqa: F401,F403 from .dinov2_arch import * # noqa: F401,F403 from .dual_hybrid_vit import * # noqa: F401,F403 from .enable_cpe_support import * # noqa: F401,F403 from .enable_spectral_reparam import * # noqa: F401,F403 from .eradio_model import * # noqa: F401,F403 from .extra_models import * # noqa: F401,F403 from .extra_timm_models import * # noqa: F401,F403 from .feature_normalizer import * # noqa: F401,F403 from .forward_intermediates import * # noqa: F401,F403 from .hf_model import RADIOConfig as HFRADIOConfig, RADIOModel as HFRADIOModel from .input_conditioner import * # noqa: F401,F403 from .open_clip_adaptor import * # noqa: F401,F403 from .radio_model import * # noqa: F401,F403 from .vit_patch_generator import * # noqa: F401,F403 from .vitdet import * # noqa: F401,F403 IGNORE_INDEX = -100 def prepare_inputs_labels_for_multimodal_vectorllm( llm, input_ids: torch.LongTensor = None, position_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, labels: Optional[torch.LongTensor] = None, pixel_values: Optional[torch.FloatTensor] = None, pixel_token_id=None, ): if pixel_values is None: return { "input_ids": input_ids, "position_ids": position_ids, "attention_mask": attention_mask, "past_key_values": past_key_values, "inputs_embeds": None, "labels": labels, } original_labels = labels original_position_ids = position_ids original_attention_mask = attention_mask if attention_mask is None: attention_mask = torch.ones_like(input_ids, dtype=torch.bool) else: attention_mask = attention_mask.bool() if position_ids is None: position_ids = torch.arange( 0, input_ids.shape[1], dtype=torch.long, device=input_ids.device ).unsqueeze(0).expand(input_ids.shape[0], -1) if labels is None: labels = torch.full_like(input_ids, IGNORE_INDEX) inputs_embeds = llm.get_input_embeddings()(input_ids) inputs_embeds = inputs_embeds.clone() labels = labels.clone() if pixel_values.ndim != 3: raise ValueError(f"Expected pixel_values to have shape [B, N, C], got {tuple(pixel_values.shape)}") for batch_idx in range(input_ids.shape[0]): replace_positions = torch.where(input_ids[batch_idx] == pixel_token_id)[0] if replace_positions.numel() == 0: continue if replace_positions.numel() != pixel_values.shape[1]: raise ValueError( "The number of image placeholder tokens does not match the projected visual tokens: " f"{replace_positions.numel()} vs {pixel_values.shape[1]}" ) inputs_embeds[batch_idx, replace_positions] = pixel_values[batch_idx].to(inputs_embeds.dtype) labels[batch_idx, replace_positions] = IGNORE_INDEX return { "input_ids": None, "position_ids": None if original_position_ids is None else position_ids, "attention_mask": None if original_attention_mask is None else attention_mask.to(dtype=original_attention_mask.dtype), "past_key_values": past_key_values, "inputs_embeds": inputs_embeds, "labels": None if original_labels is None else labels, } class ProjectorModel(PreTrainedModel): config_class = ProjectorConfig base_model_prefix = "model" supports_gradient_checkpointing = True def __init__(self, config: ProjectorConfig) -> None: super().__init__(config) self.gradient_checkpointing = False modules = [ nn.Linear(config.visual_hidden_size, config.llm_hidden_size, bias=config.bias) ] for _ in range(1, config.depth): modules.append(ACT2FN[config.hidden_act]) modules.append( nn.Linear(config.llm_hidden_size, config.llm_hidden_size, bias=config.bias) ) self.model = nn.Sequential(*modules) def forward(self, x): if self.gradient_checkpointing and self.training: return torch.utils.checkpoint.checkpoint(self.model, x) return self.model(x) class VectorLLMForCausalLM(PreTrainedModel): config_class = VectorLLMConfig main_input_name = "pixel_values" base_model_prefix = "model" supports_gradient_checkpointing = True def __init__( self, config: VectorLLMConfig, vision_model=None, language_model=None, projector=None, pos_embeds=None, ): super().__init__(config) if vision_model is not None: self.vision_model = vision_model else: self.vision_model = HFRADIOModel(HFRADIOConfig(**config.vision_config)) target_dtype = getattr(torch, config.vision_torch_dtype, None) if target_dtype is not None: self.vision_model = self.vision_model.to(dtype=target_dtype) if language_model is not None: self.language_model = language_model else: self.language_model = Qwen3ForCausalLM(Qwen3Config(**config.llm_config)) if projector is not None: self.projector = projector else: self.projector = ProjectorModel(ProjectorConfig(**config.projector_config)) width = config.regression_size[0] // config.patch_size height = config.regression_size[1] // config.patch_size n_pos = width * height if pos_embeds is not None: self.visual_pos_embeddings = pos_embeds else: self.visual_pos_embeddings = nn.Embedding(n_pos, config.vision_hidden_size) self.pixel_idx = config.pixel_idx self.num_cls_register_tokens = config.num_cls_register_tokens @property def lm_head(self): return self.language_model.get_output_embeddings() def get_input_embeddings(self): return self.language_model.get_input_embeddings() def get_output_embeddings(self): return self.language_model.get_output_embeddings() def extract_feature(self, pixel_values): summary, features = self.vision_model(pixel_values.to(self.vision_model.dtype)) del summary pos_embed = self.visual_pos_embeddings.weight.unsqueeze(0) pos_embed = pos_embed.repeat(features.shape[0], 1, 1) features = features + pos_embed features = features.to(self.projector.dtype) return self.projector(features) def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values=None, inputs_embeds: Optional[torch.FloatTensor] = 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, pixel_values: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, ): if inputs_embeds is None and pixel_values is not None: if isinstance(pixel_values, list): pixel_values = [item.unsqueeze(0) if item.ndim == 3 else item for item in pixel_values] pixel_values = torch.cat(pixel_values, dim=0) pixel_values = pixel_values.to(self.device) projected = self.extract_feature(pixel_values) llm_inputs = prepare_inputs_labels_for_multimodal_vectorllm( llm=self.language_model, input_ids=input_ids, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, labels=labels, pixel_values=projected, pixel_token_id=self.pixel_idx, ) inputs_embeds = llm_inputs["inputs_embeds"] attention_mask = llm_inputs["attention_mask"] position_ids = llm_inputs["position_ids"] labels = llm_inputs["labels"] input_ids = llm_inputs["input_ids"] outputs = self.language_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) logits = outputs.logits loss = None if labels is not None: shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_fct = CrossEntropyLoss() shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size) shift_labels = shift_labels.view(-1).to(shift_logits.device) loss = loss_fct(shift_logits, shift_labels) if not return_dict: output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) @torch.no_grad() def generate( self, pixel_values: Optional[torch.FloatTensor] = None, input_ids: Optional[torch.FloatTensor] = None, attention_mask: Optional[torch.LongTensor] = None, generation_config: Optional[GenerationConfig] = None, output_hidden_states: Optional[bool] = None, return_dict_in_generate: Optional[bool] = None, **generate_kwargs, ) -> torch.LongTensor: if pixel_values is not None: if isinstance(pixel_values, list): pixel_values = [item.unsqueeze(0) if item.ndim == 3 else item for item in pixel_values] pixel_values = torch.cat(pixel_values, dim=0) pixel_values = pixel_values.to(self.device) input_ids = input_ids.to(self.device) input_embeds = self.language_model.get_input_embeddings()(input_ids) projected = self.extract_feature(pixel_values).to(input_embeds.dtype) batch, seqlen, channels = input_embeds.shape flat_embeds = input_embeds.reshape(batch * seqlen, channels) selected = input_ids.reshape(batch * seqlen) == self.pixel_idx flat_embeds[selected] = projected.reshape(-1, channels).to(flat_embeds.device) input_embeds = flat_embeds.reshape(batch, seqlen, channels) else: input_embeds = self.language_model.get_input_embeddings()(input_ids.to(self.device)) outputs = self.language_model.generate( inputs_embeds=input_embeds, attention_mask=attention_mask.to(self.device) if attention_mask is not None else None, generation_config=generation_config, output_hidden_states=output_hidden_states, return_dict_in_generate=return_dict_in_generate, **generate_kwargs, ) return outputs