# -*- coding: utf-8 -*- """RWKV-VL HF remote-code model used by the exporter. This file intentionally lives beside the exporter so the generated checkpoint is self-contained. """ from dataclasses import dataclass import warnings from typing import Any, Dict, Optional, Tuple, Union import torch import torch.nn as nn from transformers.cache_utils import Cache from transformers import ( AutoConfig, AutoModel, AutoModelForCausalLM, AutoModelForImageTextToText, PretrainedConfig, PreTrainedModel, Qwen3VLVisionModel, ) from transformers.generation import GenerationMixin from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLVisionConfig try: from .configuration_rwkv7 import RWKV7Config from .modeling_rwkv7 import RWKV7Model except ImportError: from configuration_rwkv7 import RWKV7Config from modeling_rwkv7 import RWKV7Model @dataclass class ModRWKVProjectorConfig: projector_type: str = "visual" encoder_dim: int = 1024 project_dim: int = 1024 hidden_dim: Optional[int] = None num_deepstack: int = 0 def to_dict(self) -> dict[str, Any]: return { "projector_type": self.projector_type, "encoder_dim": self.encoder_dim, "project_dim": self.project_dim, "hidden_dim": self.hidden_dim, "num_deepstack": self.num_deepstack, } class ModRWKVConfig(PretrainedConfig): model_type = "modrwkv" is_composition = True @staticmethod def _to_vision_config( vision_config: Union[Qwen3VLVisionConfig, Dict[str, Any]], ) -> Qwen3VLVisionConfig: if isinstance(vision_config, Qwen3VLVisionConfig): return vision_config if not isinstance(vision_config, dict) and hasattr(vision_config, "to_dict"): vision_config = vision_config.to_dict() if isinstance(vision_config, dict): if not vision_config: return Qwen3VLVisionConfig() model_type = vision_config.get("model_type") if model_type not in {None, Qwen3VLVisionConfig.model_type, "qwen3_vl"}: raise TypeError( "ModRWKVConfig expects a Qwen3-VL vision config; " f"got model_type={model_type!r}." ) return Qwen3VLVisionConfig( depth=vision_config["depth"], hidden_size=vision_config["hidden_size"], hidden_act=vision_config.get("hidden_act", "gelu_pytorch_tanh"), intermediate_size=vision_config["intermediate_size"], num_heads=vision_config["num_heads"], in_channels=vision_config.get("in_channels", 3), patch_size=vision_config.get("patch_size", 16), spatial_merge_size=vision_config.get("spatial_merge_size", 2), temporal_patch_size=vision_config.get("temporal_patch_size", 2), out_hidden_size=vision_config["out_hidden_size"], num_position_embeddings=vision_config.get( "num_position_embeddings", 2304, ), deepstack_visual_indexes=list( vision_config.get("deepstack_visual_indexes") or vision_config.get("deepstack_visual_indices") or [] ), initializer_range=vision_config.get("initializer_range", 0.02), ) raise TypeError(f"Unsupported vision config type: {type(vision_config)!r}") @classmethod def from_text_vision_configs( cls, text_config: Union[RWKV7Config, Dict[str, Any]], vision_config: Union[Qwen3VLVisionConfig, Dict[str, Any]], projector_config: Optional[Union[ModRWKVProjectorConfig, Dict[str, Any]]] = None, **kwargs, ) -> "ModRWKVConfig": return cls( text_config=text_config, vision_config=vision_config, projector_config=projector_config, **kwargs, ) def __init__( self, text_config: Optional[Union[RWKV7Config, Dict[str, Any]]] = None, vision_config: Optional[Union[Qwen3VLVisionConfig, Dict[str, Any]]] = None, projector_config: Optional[Union[ModRWKVProjectorConfig, Dict[str, Any]]] = None, image_token_id: int = 65532, vision_start_token_id: int = 65530, vision_end_token_id: int = 65531, tie_word_embeddings: bool = False, use_conv_in_projector: bool = False, **kwargs, ): super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) if text_config is None: text_config = {} if isinstance(text_config, dict): text_config = RWKV7Config(**text_config) self.text_config = text_config if vision_config is None: vision_config = {} vision_config = self._to_vision_config(vision_config) self.vision_config = vision_config if projector_config is None: deepstack_indexes = getattr( vision_config, "deepstack_visual_indexes", getattr(vision_config, "deepstack_visual_indices", []), ) projector_config = ModRWKVProjectorConfig( encoder_dim=getattr(vision_config, "out_hidden_size", 1024), project_dim=getattr(text_config, "hidden_size", 1024), num_deepstack=len(deepstack_indexes), ) elif isinstance(projector_config, dict): projector_config = ModRWKVProjectorConfig(**projector_config) self.projector_config = projector_config self.image_token_id = image_token_id self.vision_start_token_id = vision_start_token_id self.vision_end_token_id = vision_end_token_id self.use_conv_in_projector = use_conv_in_projector def to_dict(self) -> dict[str, Any]: output = super().to_dict() output["text_config"] = ( self.text_config.to_dict() if hasattr(self.text_config, "to_dict") else self.text_config ) output["vision_config"] = ( self.vision_config.to_dict() if hasattr(self.vision_config, "to_dict") else self.vision_config ) output["projector_config"] = ( self.projector_config.to_dict() if hasattr(self.projector_config, "to_dict") else self.projector_config ) output["image_token_id"] = self.image_token_id output["vision_start_token_id"] = self.vision_start_token_id output["vision_end_token_id"] = self.vision_end_token_id output["use_conv_in_projector"] = self.use_conv_in_projector return output class ModRWKVPreTrainedModel(PreTrainedModel): config_class = ModRWKVConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["RWKV7Block"] _supports_cache_class = True _skip_keys_device_placement = ["past_key_values"] class _VisualStreamProjector(nn.Module): def __init__( self, encoder_dim: int, project_dim: int, hidden_dim: Optional[int] = None, ): super().__init__() self.encoder_dim = encoder_dim self.project_dim = project_dim self.hidden_dim = hidden_dim or project_dim * 4 self.pre_norm = nn.LayerNorm(project_dim) self.mlp = nn.Sequential( nn.Linear(encoder_dim, self.hidden_dim), nn.ReLU(), nn.Linear(self.hidden_dim, project_dim), ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.mlp(x) return x + self.pre_norm(x) class VisualAdapter(nn.Module): def __init__( self, encoder_dim: int, project_dim: int, hidden_dim: Optional[int] = None, num_deepstack: int = 0, use_conv: bool = False, ): super().__init__() if use_conv: raise ValueError("Convolutional visual projectors are not supported.") self.encoder_dim = encoder_dim self.project_dim = project_dim self.hidden_dim = hidden_dim or project_dim * 4 self.num_deepstack = num_deepstack self.main = _VisualStreamProjector( encoder_dim=encoder_dim, project_dim=project_dim, hidden_dim=self.hidden_dim, ) self.deepstack = nn.ModuleList( [ _VisualStreamProjector( encoder_dim=encoder_dim, project_dim=project_dim, hidden_dim=self.hidden_dim, ) for _ in range(num_deepstack) ] ) def forward( self, x: torch.Tensor, deepstack_features: Optional[list[torch.Tensor]] = None, ) -> tuple[torch.Tensor, list[torch.Tensor]]: if deepstack_features is None: deepstack_features = [] if len(deepstack_features) != self.num_deepstack: raise ValueError( f"Expected {self.num_deepstack} DeepStack feature tensors, " f"got {len(deepstack_features)}." ) projected_deepstack = [ projector(feature) for projector, feature in zip(self.deepstack, deepstack_features) ] return self.main(x), projected_deepstack class RWKV7VLModel(ModRWKVPreTrainedModel): def __init__(self, config: ModRWKVConfig): super().__init__(config) self.encoder = Qwen3VLVisionModel(config.vision_config) proj_cfg = config.projector_config self.proj = VisualAdapter( encoder_dim=proj_cfg.encoder_dim, project_dim=proj_cfg.project_dim, hidden_dim=proj_cfg.hidden_dim, num_deepstack=proj_cfg.num_deepstack, use_conv=config.use_conv_in_projector, ) self.llm = RWKV7Model(config.text_config) self.post_init() def get_input_embeddings(self): return self.llm.get_input_embeddings() def set_input_embeddings(self, value): self.llm.set_input_embeddings(value) def _get_image_features( self, pixel_values: torch.FloatTensor, image_grid_thw: torch.LongTensor, ) -> tuple[torch.FloatTensor, list[torch.FloatTensor]]: vision_output = self.encoder(pixel_values, image_grid_thw) if hasattr(vision_output, "pooler_output"): vision_embeds = vision_output.pooler_output elif hasattr(vision_output, "last_hidden_state"): vision_embeds = vision_output.last_hidden_state elif isinstance(vision_output, (tuple, list)): vision_embeds = vision_output[0] else: vision_embeds = vision_output deepstack_features = getattr(vision_output, "deepstack_features", None) if deepstack_features is None: deepstack_features = [] projected, projected_deepstack = self.proj( vision_embeds, list(deepstack_features), ) projected = projected.reshape(-1, self.config.text_config.hidden_size) projected_deepstack = [ feature.reshape(-1, self.config.text_config.hidden_size) for feature in projected_deepstack ] spatial_merge_size = getattr(self.encoder.config, "spatial_merge_size", 2) expected_tokens = int( (image_grid_thw.prod(-1) // (spatial_merge_size**2)).sum().item() ) if expected_tokens != projected.shape[0]: raise ValueError( "Projected image features and image grid do not match: " f"features={projected.shape[0]} grid_tokens={expected_tokens}" ) if projected.numel() == 0: empty = torch.empty( 0, self.config.text_config.hidden_size, device=self.get_input_embeddings().weight.device, ) return empty, [] return projected, projected_deepstack def _inject_image_features( self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor, ) -> torch.FloatTensor: image_mask = input_ids == self.config.image_token_id if image_mask.sum().item() != image_features.shape[0]: raise ValueError( "Image features and image placeholder tokens do not match: " f"tokens={image_mask.sum().item()} features={image_features.shape[0]}" ) inputs_embeds = inputs_embeds.clone() inputs_embeds[image_mask] = image_features.to( device=inputs_embeds.device, dtype=inputs_embeds.dtype, ) return inputs_embeds def _add_image_features( self, input_ids: torch.LongTensor, hidden_states: torch.FloatTensor, image_features: torch.FloatTensor, ) -> torch.FloatTensor: image_mask = input_ids == self.config.image_token_id if image_mask.sum().item() != image_features.shape[0]: raise ValueError( "DeepStack features and image placeholder tokens do not match: " f"tokens={image_mask.sum().item()} features={image_features.shape[0]}" ) hidden_states = hidden_states.clone() hidden_states[image_mask] += image_features.to( device=hidden_states.device, dtype=hidden_states.dtype, ) return hidden_states def forward( self, input_ids: Optional[torch.LongTensor] = None, pixel_values: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, past_key_values: Optional[Any] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, **kwargs, ) -> Union[Tuple, BaseModelOutputWithPast]: return_dict = ( return_dict if return_dict is not None else self.config.text_config.use_return_dict ) output_attentions = ( output_attentions if output_attentions is not None else self.config.text_config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.text_config.output_hidden_states ) use_cache = ( use_cache if use_cache is not None else (self.config.text_config.use_cache if not self.training else False) ) if output_attentions: warnings.warn( "`RWKV7Model` does not support `output_attentions`; setting it to `False`." ) output_attentions = False if input_ids is None and inputs_embeds is None: raise ValueError("You must provide either input_ids or inputs_embeds.") if (pixel_values is None) != (image_grid_thw is None): raise ValueError("pixel_values and image_grid_thw must be provided together.") if pixel_values is not None and input_ids is None: raise ValueError("input_ids are required when pixel_values are provided.") if inputs_embeds is None: inputs_embeds = self.get_input_embeddings()(input_ids) deepstack_features: list[torch.Tensor] = [] if pixel_values is not None: image_features, deepstack_features = self._get_image_features( pixel_values, image_grid_thw, ) inputs_embeds = self._inject_image_features( input_ids, inputs_embeds, image_features, ) if use_cache and past_key_values is not None and not isinstance( past_key_values, Cache, ): from_legacy_cache = getattr(Cache, "from_legacy_cache", None) if callable(from_legacy_cache): past_key_values = from_legacy_cache(past_key_values) all_hidden_states = () if output_hidden_states else None all_attns = () if output_attentions else None hidden_states = inputs_embeds v_first = torch.zeros_like(hidden_states) for layer_idx, layer in enumerate(self.llm.layers): if output_hidden_states: all_hidden_states += (hidden_states,) hidden_states, attentions, past_key_values, v_first = layer( hidden_states, attention_mask=attention_mask, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, v_first=v_first, **kwargs, ) if layer_idx < len(deepstack_features): hidden_states = self._add_image_features( input_ids, hidden_states, deepstack_features[layer_idx], ) if output_attentions: all_attns += (attentions,) hidden_states = self.llm.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) if not return_dict: return tuple( item for item in [ hidden_states, past_key_values, all_hidden_states, all_attns, ] if item is not None ) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, hidden_states=all_hidden_states, attentions=all_attns, ) class RWKV7VLForConditionalGeneration(ModRWKVPreTrainedModel, GenerationMixin): _tied_weights_keys = {} def __init__(self, config: ModRWKVConfig): super().__init__(config) self.model = RWKV7VLModel(config) self.vocab_size = config.text_config.vocab_size self.lm_head = nn.Linear( config.text_config.hidden_size, config.text_config.vocab_size, bias=False, ) self.post_init() def get_input_embeddings(self): return self.model.get_input_embeddings() def set_input_embeddings(self, value): self.model.set_input_embeddings(value) def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings @staticmethod def _has_nonempty_past_key_values( past_key_values: Optional[Any], cache_position: Optional[torch.LongTensor], ) -> bool: if past_key_values is None: return False if cache_position is not None: return cache_position.numel() > 0 and cache_position[0].item() > 0 get_seq_length = getattr(past_key_values, "get_seq_length", None) if callable(get_seq_length): try: return get_seq_length() > 0 except (AttributeError, TypeError): pass try: return len(past_key_values) > 0 except TypeError: return True def prepare_inputs_for_generation( self, input_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Any] = None, attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, pixel_values: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = True, logits_to_keep: Optional[int] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs, ) -> dict[str, Any]: has_past = self._has_nonempty_past_key_values( past_key_values, cache_position, ) if has_past and input_ids is not None: if ( cache_position is not None and input_ids.shape[1] != cache_position.shape[0] ): input_ids = input_ids[:, cache_position] else: input_ids = input_ids[:, -1:] model_inputs: dict[str, Any] = { "input_ids": input_ids.contiguous() if input_ids is not None else None, "past_key_values": past_key_values, "attention_mask": attention_mask, "use_cache": use_cache, } if inputs_embeds is not None and not has_past: model_inputs["inputs_embeds"] = inputs_embeds if not has_past: model_inputs["pixel_values"] = pixel_values model_inputs["image_grid_thw"] = image_grid_thw if cache_position is not None: model_inputs["cache_position"] = cache_position if logits_to_keep is not None: model_inputs["logits_to_keep"] = logits_to_keep return model_inputs def forward( self, input_ids: Optional[torch.LongTensor] = None, pixel_values: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, past_key_values: Optional[Any] = None, labels: Optional[torch.LongTensor] = None, shift_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, logits_to_keep: Optional[int] = 0, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: return_dict = ( return_dict if return_dict is not None else self.config.text_config.use_return_dict ) outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, image_grid_thw=image_grid_thw, attention_mask=attention_mask, inputs_embeds=inputs_embeds, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, **kwargs, ) hidden_states = outputs.last_hidden_state logits = self.lm_head( hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:] ) loss = None if labels is not None or shift_labels is not None: if shift_labels is None: ignore = torch.full_like(labels[:, :1], -100) shift_labels = torch.cat((labels[..., 1:], ignore), dim=1) loss = nn.CrossEntropyLoss()( logits.reshape(-1, logits.shape[-1]), shift_labels.to(logits.device).reshape(-1), ) 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, ) AutoConfig.register(ModRWKVConfig.model_type, ModRWKVConfig, exist_ok=True) AutoModel.register(ModRWKVConfig, RWKV7VLForConditionalGeneration, exist_ok=True) AutoModelForCausalLM.register(ModRWKVConfig, RWKV7VLForConditionalGeneration, exist_ok=True) AutoModelForImageTextToText.register( ModRWKVConfig, RWKV7VLForConditionalGeneration, exist_ok=True, ) ModRWKVConfig.register_for_auto_class("AutoConfig") RWKV7VLForConditionalGeneration.register_for_auto_class("AutoModel") RWKV7VLForConditionalGeneration.register_for_auto_class("AutoModelForCausalLM") RWKV7VLForConditionalGeneration.register_for_auto_class("AutoModelForImageTextToText")