Add TIPSv2 vision-language dual-encoder model to hub#2721
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the TIPSv2 model architecture, featuring a dual-encoder backbone with DINOv2-style vision and custom text encoders, along with preprocessors, tokenizers, and checkpoint conversion utilities. Feedback identifies that the contrastive temperature should be a trainable weight and is currently omitted from the model's computation graph. Furthermore, the vision encoder's positional embedding interpolation must be made dynamic to support arbitrary resolutions, and the image converter's default size needs to be corrected to 448x448. Finally, duplicate implementations of the sinusoidal position embedding layer should be merged into a single efficient shared component.
| ) | ||
|
|
||
| # === Config === | ||
| self.temperature = temperature |
There was a problem hiding this comment.
The temperature parameter is described as a "learned contrastive temperature" in the docstring and PR description, but it is currently implemented as a fixed Python attribute. In contrastive dual-encoder models, this is typically a trainable weight (often log-scaled) used to scale the similarity scores. Furthermore, it is not used anywhere in the model's functional graph (the call method), meaning it has no effect on the model's outputs or training. If it is intended to be learned, it should be initialized as a trainable Variable and applied to the embeddings or output as part of the model's results.
| w0 = self.image_size // self.patch_size | ||
| h0 = self.image_size // self.patch_size |
There was a problem hiding this comment.
The _interpolate_pos_encoding method uses self.image_size (a fixed value from initialization) to determine the target grid size for interpolation. This breaks the "interpolation support for arbitrary resolutions" mentioned in the class docstring. If an input image with a different resolution is provided, the interpolated positional embeddings will have a shape mismatch with the patch embeddings. The target grid size (h0, w0) should be calculated dynamically based on the input tensor's spatial dimensions.
| class TIPSv2ImageConverter(ImageConverter): | ||
| """TIPSv2 image converter. | ||
|
|
||
| Resizes images and scales pixel values to [0, 1]. | ||
| Unlike CLIP/SigLIP, TIPSv2 does NOT apply ImageNet normalization. | ||
|
|
||
| The default image size is 448x448 as used in all TIPSv2 variants. | ||
| """ | ||
|
|
||
| backbone_cls = TIPSv2Backbone |
There was a problem hiding this comment.
The docstring states that the default image size is 448x448, but the class does not override the __init__ method to set this default. Consequently, it will use the default from the ImageConverter base class (typically 224x224), which contradicts the documentation and the model's expected input resolution. Please add an __init__ method to set the correct default image_size.
| class TIPSv2ImageConverter(ImageConverter): | |
| """TIPSv2 image converter. | |
| Resizes images and scales pixel values to [0, 1]. | |
| Unlike CLIP/SigLIP, TIPSv2 does NOT apply ImageNet normalization. | |
| The default image size is 448x448 as used in all TIPSv2 variants. | |
| """ | |
| backbone_cls = TIPSv2Backbone | |
| @keras_hub_export("keras_hub.layers.TIPSv2ImageConverter") | |
| class TIPSv2ImageConverter(ImageConverter): | |
| """TIPSv2 image converter. | |
| Resizes images and scales pixel values to [0, 1]. | |
| Unlike CLIP/SigLIP, TIPSv2 does NOT apply ImageNet normalization. | |
| The default image size is 448x448 as used in all TIPSv2 variants. | |
| """ | |
| backbone_cls = TIPSv2Backbone | |
| def __init__(self, image_size=(448, 448), **kwargs): | |
| super().__init__(image_size=image_size, **kwargs) |
| class TIPSv2SinusoidalPositionEmbedding(keras.layers.Layer): | ||
| """Sinusoidal positional embedding (computed, not learned). | ||
|
|
||
| Generates position embeddings using sin/cos functions with geometric | ||
| timescale progression, matching the TIPSv2 text encoder reference. | ||
|
|
||
| Args: | ||
| embedding_dim: int. Dimension of the positional embedding. | ||
| min_timescale: int. Minimum timescale. Defaults to `1`. | ||
| max_timescale: int. Maximum timescale. Defaults to `10000`. | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| embedding_dim, | ||
| min_timescale=1, | ||
| max_timescale=10000, | ||
| **kwargs, | ||
| ): | ||
| super().__init__(**kwargs) | ||
| self.embedding_dim = embedding_dim | ||
| self.min_timescale = min_timescale | ||
| self.max_timescale = max_timescale | ||
|
|
||
| def call(self, seq_length): | ||
| num_timescales = self.embedding_dim // 2 | ||
| log_timescale_increment = math.log( | ||
| float(self.max_timescale) / float(self.min_timescale) | ||
| ) / max(num_timescales - 1, 1) | ||
|
|
||
| inv_timescales = self.min_timescale * ops.exp( | ||
| ops.cast(ops.arange(num_timescales), "float32") | ||
| * (-log_timescale_increment) | ||
| ) | ||
| # positions: (1, seq_length, 1) | ||
| position = ops.cast(ops.arange(seq_length), "float32")[None, :, None] | ||
| # inv_timescales: (1, 1, num_timescales) | ||
| inv_timescales = inv_timescales[None, None, :] | ||
|
|
||
| scaled_time = position * inv_timescales # (1, seq, num_ts) | ||
| signal = ops.concatenate( | ||
| [ops.sin(scaled_time), ops.cos(scaled_time)], axis=-1 | ||
| ) # (1, seq, embedding_dim) or (1, seq, embedding_dim-1) | ||
|
|
||
| # Pad if embedding_dim is odd. | ||
| if self.embedding_dim % 2 != 0: | ||
| signal = ops.pad(signal, [[0, 0], [0, 0], [0, 1]]) | ||
|
|
||
| return signal | ||
|
|
||
| def get_config(self): | ||
| config = super().get_config() | ||
| config.update( | ||
| { | ||
| "embedding_dim": self.embedding_dim, | ||
| "min_timescale": self.min_timescale, | ||
| "max_timescale": self.max_timescale, | ||
| } | ||
| ) | ||
| return config |
There was a problem hiding this comment.
This class is a duplicate of the one defined in tipsv2_text_encoder.py. The implementation here is also less efficient because it recomputes the sinusoidal signal on every call, whereas the version in tipsv2_text_encoder.py correctly precomputes it in build() and stores it as a non-trainable weight. You should remove this duplicate and ensure the efficient version is used consistently (ideally by keeping it in this shared layers file and importing it where needed).
| class TIPSv2SinusoidalPositionEmbedding(keras.layers.Layer): | ||
| """Sinusoidal positional embedding (computed, not learned). | ||
|
|
||
| Takes an input tensor and returns a position embedding tensor of the | ||
| same shape, computed using sin/cos functions. This layer broadcasts | ||
| the embedding to match the input's batch and sequence dimensions. | ||
|
|
||
| Args: | ||
| max_sequence_length: int. Maximum supported sequence length. | ||
| embedding_dim: int. Dimension of the positional embedding. | ||
| min_timescale: int. Minimum timescale. Defaults to `1`. | ||
| max_timescale: int. Maximum timescale. Defaults to `10000`. | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| max_sequence_length, | ||
| embedding_dim, | ||
| min_timescale=1, | ||
| max_timescale=10000, | ||
| **kwargs, | ||
| ): | ||
| super().__init__(**kwargs) | ||
| self.max_sequence_length = max_sequence_length | ||
| self.embedding_dim = embedding_dim | ||
| self.min_timescale = min_timescale | ||
| self.max_timescale = max_timescale | ||
|
|
||
| def build(self, input_shape): | ||
| # Pre-compute the full position embedding table at build time. | ||
| num_timescales = self.embedding_dim // 2 | ||
| log_timescale_increment = math.log( | ||
| float(self.max_timescale) / float(self.min_timescale) | ||
| ) / max(num_timescales - 1, 1) | ||
|
|
||
| inv_timescales = self.min_timescale * ops.exp( | ||
| ops.cast(ops.arange(num_timescales), "float32") | ||
| * (-log_timescale_increment) | ||
| ) | ||
| position = ops.cast(ops.arange(self.max_sequence_length), "float32")[ | ||
| :, None | ||
| ] | ||
| inv_timescales = inv_timescales[None, :] | ||
|
|
||
| scaled_time = position * inv_timescales # (max_len, num_ts) | ||
| signal = ops.concatenate( | ||
| [ops.sin(scaled_time), ops.cos(scaled_time)], axis=-1 | ||
| ) # (max_len, embedding_dim) | ||
|
|
||
| # Pad if embedding_dim is odd. | ||
| if self.embedding_dim % 2 != 0: | ||
| signal = ops.pad(signal, [[0, 0], [0, 1]]) | ||
|
|
||
| # Store as non-trainable weight for serialization. | ||
| self.pos_table = self.add_weight( | ||
| name="pos_table", | ||
| shape=(self.max_sequence_length, self.embedding_dim), | ||
| initializer="zeros", | ||
| trainable=False, | ||
| ) | ||
| self.pos_table.assign(signal) | ||
| self.built = True | ||
|
|
||
| def call(self, inputs): | ||
| """Return position embeddings matching input sequence length. | ||
|
|
||
| Args: | ||
| inputs: Tensor of shape (B, seq_len, D). | ||
|
|
||
| Returns: | ||
| Tensor of shape (B, seq_len, D) with position embeddings. | ||
| """ | ||
| seq_length = ops.shape(inputs)[-2] | ||
| # Slice the pre-computed table to match seq_length. | ||
| # Explicitly convert Variable to tensor for JAX tracing. | ||
| pos_table = ops.convert_to_tensor(self.pos_table) | ||
| pos_embed = ops.slice( | ||
| pos_table, (0, 0), (seq_length, self.embedding_dim) | ||
| ) | ||
| return ops.broadcast_to(pos_embed, ops.shape(inputs)) | ||
|
|
||
| def compute_output_shape(self, input_shape): | ||
| return input_shape | ||
|
|
||
| def get_config(self): | ||
| config = super().get_config() | ||
| config.update( | ||
| { | ||
| "max_sequence_length": self.max_sequence_length, | ||
| "embedding_dim": self.embedding_dim, | ||
| "min_timescale": self.min_timescale, | ||
| "max_timescale": self.max_timescale, | ||
| } | ||
| ) | ||
| return config |
There was a problem hiding this comment.
This class is a duplicate of the one defined in tipsv2_layers.py. While this version is more efficient (it precomputes the table), having two identical classes in the same model directory is bad for maintainability. You should move this efficient implementation to tipsv2_layers.py and import it here.
laxmareddyp
added
the
new model
1 participant
Description of the change
This PR solves the issue #2703
Summary:
Paper: TIPSv2: Advancing Vision-Language Pretraining with Enhanced Patch-Text Alignment (CVPR 2026)
Project Page: https://gdm-tipsv2.github.io/
Source: https://github.com/google-deepmind/tips
Architecture Overview
TIPSv2 is a dual-encoder architecture with separate vision and text encoders that map inputs into a shared embedding space:
1. Vision Encoder — A DINOv2-style Vision Transformer with:
2.Text Encoder — A custom transformer encoder with:
3.Backbone — Wraps both encoders into a single functional model with a learned contrastive temperature parameter. Produces aligned vision_cls_embedding, vision_patch_embeddings, and text_embedding outputs.
References
Paper: TIPSv2: Advancing Vision-Language Pretraining with Enhanced Patch-Text Alignment (CVPR 2026)
Project Page: https://gdm-tipsv2.github.io/
GitHub: https://github.com/google-deepmind/tips
HuggingFace Models: https://huggingface.co/collections/google/tipsv2
HF Demo: https://huggingface.co/spaces/google/TIPSv2
Architecture Base: DINOv2 ViT + Custom Text Encoder with Sinusoidal PE + Contrastive Learning