Add blip2 (keras-team#2699)

* Add BLIP-2 and InstructBLIP models

* Split BLIP-2 Flan-T5 into a BLIP2Seq2SeqLM task

The Flan-T5 variant is encoder-decoder but was running through
BLIP2CausalLM via is_encoder_decoder branching, which also double-encoded
the image (vision encoder + Q-Former ran once in _encode_images() and
again directly in generate_step).

Add BLIP2Seq2SeqLM(Seq2SeqLM) and BLIP2Seq2SeqLMPreprocessor following the
T5Gemma(2)Seq2SeqLM convention: encoder_token_ids/decoder_token_ids inputs,
a single vision/Q-Former encode in call_encoder, and a clean seq2seq
generate_step. BLIP2CausalLM is now strictly decoder-only (OPT/Vicuna), and
BLIP2Backbone exposes encoder_*/padding_mask names for the encoder-decoder
LM. T5 has no KV cache, so the decoder recomputes each step.

* Package BLIP-2 Flan-T5 presets as BLIP2Seq2SeqLM

* Resolving Maintainer comments

* run pre commit hook

* Resolving maintainer comments-2

* Remove the InstructBLIP-Vicuna variant from BLIP-2

* Run the conversion tool on GPU instead of forcing CPU

* Keep TensorFlow on CPU during BLIP-2 checkpoint conversion

The conversion runs on the torch backend, but keras-hub tokenizers use
TensorFlow string ops. On newer GPUs (e.g. Blackwell / compute
capability 12.0) some TF builds have no matching CUDA kernels and fail
to JIT them from PTX, crashing the OPT BPE tokenizer's tf.cast with
CUDA_ERROR_INVALID_PTX. Hide the GPU from TensorFlow so its tokenizer
string ops stay on CPU while torch/Keras keep the GPU for the model.

* force fp32 matmuls in BLIP-2 conversion

disable TF32 for GPU matmuls/convs so the parity check against the
float32 HuggingFace model stays tight (~1e-4) instead of drifting to
~1e-3 from TF32's reduced mantissa.

* fix pytorch and tensorflow kernel kill issue

* Mark heavyweight BLIP-2 tests as extra_large to fix GPU CI OOM

The class detection tests build the full blip2-opt-2.7b and
blip2-flan-t5-xl architectures (~4B params each, twice per test), which
exhausts the 16GB GPU in the large test job and cascades OOM and
corrupted Keras stateless-scope state into every test that runs after
them. Also move test_xl_config_instantiates (~2.8B params) out of the
large job since it bloats the TF allocator pool for the whole session.

* Mark BLIP-2 saved model tests as large per repo convention

The save/load round-trip tests in the backbone, causal LM, and seq2seq
LM test files were unmarked, so they ran in the small presubmit. Mark
them @pytest.mark.large to match how other models (Gemma, Mistral) and
the other BLIP-2 test files gate these slow tests.

---------

Co-authored-by: yassienashrafwasfy <21-101112@students.eui.edu.eg>

Author: yassienashrafwasfy

keras_hub/api/layers/__init__.py

@@ -72,6 +72,9 @@
 from keras_hub.src.models.basnet.basnet_image_converter import (
     BASNetImageConverter as BASNetImageConverter,
 )
+from keras_hub.src.models.blip2.blip2_image_converter import (
+    BLIP2ImageConverter as BLIP2ImageConverter,
+)
 from keras_hub.src.models.clip.clip_image_converter import (
     CLIPImageConverter as CLIPImageConverter,
 )

keras_hub/api/models/__init__.py

@@ -76,6 +76,39 @@
 from keras_hub.src.models.bert.bert_tokenizer import (
     BertTokenizer as BertTokenizer,
 )
+from keras_hub.src.models.blip2.blip2_backbone import (
+    BLIP2Backbone as BLIP2Backbone,
+)
+from keras_hub.src.models.blip2.blip2_causal_lm import (
+    BLIP2CausalLM as BLIP2CausalLM,
+)
+from keras_hub.src.models.blip2.blip2_causal_lm_preprocessor import (
+    BLIP2CausalLMPreprocessor as BLIP2CausalLMPreprocessor,
+)
+from keras_hub.src.models.blip2.blip2_custom_opt import (
+    BLIP2CustomOPT as BLIP2CustomOPT,
+)
+from keras_hub.src.models.blip2.blip2_flan_t5_lm import (
+    BLIP2FlanT5 as BLIP2FlanT5,
+)
+from keras_hub.src.models.blip2.blip2_flan_t5_tokenizer import (
+    BLIP2FlanT5Tokenizer as BLIP2FlanT5Tokenizer,
+)
+from keras_hub.src.models.blip2.blip2_opt_tokenizer import (
+    BLIP2OPTTokenizer as BLIP2OPTTokenizer,
+)
+from keras_hub.src.models.blip2.blip2_qformer import (
+    BLIP2QFormer as BLIP2QFormer,
+)
+from keras_hub.src.models.blip2.blip2_seq_2_seq_lm import (
+    BLIP2Seq2SeqLM as BLIP2Seq2SeqLM,
+)
+from keras_hub.src.models.blip2.blip2_seq_2_seq_lm_preprocessor import (
+    BLIP2Seq2SeqLMPreprocessor as BLIP2Seq2SeqLMPreprocessor,
+)
+from keras_hub.src.models.blip2.blip2_vision_encoder import (
+    BLIP2VisionEncoder as BLIP2VisionEncoder,
+)
 from keras_hub.src.models.bloom.bloom_backbone import (
     BloomBackbone as BloomBackbone,
 )

keras_hub/api/tokenizers/__init__.py

@@ -13,6 +13,12 @@
 from keras_hub.src.models.bert.bert_tokenizer import (
     BertTokenizer as BertTokenizer,
 )
+from keras_hub.src.models.blip2.blip2_flan_t5_tokenizer import (
+    BLIP2FlanT5Tokenizer as BLIP2FlanT5Tokenizer,
+)
+from keras_hub.src.models.blip2.blip2_opt_tokenizer import (
+    BLIP2OPTTokenizer as BLIP2OPTTokenizer,
+)
 from keras_hub.src.models.bloom.bloom_tokenizer import (
     BloomTokenizer as BloomTokenizer,
 )

keras_hub/src/models/blip2/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.blip2.blip2_backbone import BLIP2Backbone
+from keras_hub.src.models.blip2.blip2_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, BLIP2Backbone)

keras_hub/src/models/blip2/blip2_backbone.py

@@ -0,0 +1,202 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.backbone import Backbone
+
+
+@keras_hub_export("keras_hub.models.BLIP2Backbone")
+class BLIP2Backbone(Backbone):
+    """BLIP-2 core network.
+
+    BLIP-2 is a vision-language model that connects a frozen image encoder
+    and a frozen large language model (LLM) through a lightweight trainable
+    Querying Transformer (Q-Former). The Q-Former distills visual information
+    into a fixed number of query embeddings which are then fed as a soft visual
+    prompt to the language model.
+
+    The forward pass follows three stages:
+      1. A `BLIP2VisionEncoder` (ViT) maps raw images to patch features.
+      2. A `BLIP2QFormer` cross-attends learned query tokens against those
+         patch features and produces a compact set of visual embeddings.
+      3. A language model (OPT or Flan-T5) receives the query embeddings
+         prepended to its token sequence and generates text.
+
+    When `vision_encoder` is `None` the backbone operates in text-only mode:
+    the Q-Former is bypassed and the language model receives only token ids.
+
+    For a higher-level text-generation interface see
+    `keras_hub.models.BLIP2CausalLM`.
+
+    Args:
+        vision_encoder: A `keras_hub.models.BLIP2VisionEncoder` instance. Pass
+            `None` for a text-only backbone.
+        qformer: A `keras_hub.models.BLIP2QFormer` instance. Pass `None` when
+            `vision_encoder` is `None`.
+        language_model: The language model instance (e.g. `BLIP2CustomOPT`
+            or `BLIP2FlanT5`).
+        dtype: string or `keras.mixed_precision.DTypePolicy`. Dtype used for
+            model computations and weights. Defaults to `None` (Keras global
+            default).
+        **kwargs: Additional keyword arguments forwarded to the base
+            `Backbone` / `keras.Model`.
+
+    Example:
+    ```python
+    # Text-only (no vision encoder)
+    backbone = keras_hub.models.BLIP2Backbone(
+        vision_encoder=None,
+        qformer=None,
+        language_model=my_opt,
+    )
+    output = backbone({"token_ids": token_ids, "padding_mask": mask})
+
+    # Full vision-language backbone
+    backbone = keras_hub.models.BLIP2Backbone.from_preset("blip2_opt_2.7b")
+    output = backbone({
+        "images": images,           # (B, H, W, 3)
+        "token_ids": token_ids,     # (B, seq_len)
+        "padding_mask": mask,       # (B, seq_len)
+    })
+    ```
+    """
+
+    def __init__(
+        self,
+        vision_encoder,
+        qformer,
+        language_model,
+        dtype=None,
+        **kwargs,
+    ):
+        self.vision_encoder = vision_encoder
+        self.qformer = qformer
+        self.language_model = language_model
+
+        multimodal = self.vision_encoder is not None
+
+        # Encoder-decoder language models (Flan-T5) follow the keras-hub
+        # seq2seq convention and name their encoder inputs
+        # `encoder_token_ids` / `encoder_padding_mask`; decoder-only language
+        # models (OPT) keep the plain `token_ids` / `padding_mask`.
+        is_encoder_decoder = hasattr(
+            self.language_model, "encoder_transformer_layers"
+        )
+        if is_encoder_decoder:
+            token_ids_name, padding_mask_name = (
+                "encoder_token_ids",
+                "encoder_padding_mask",
+            )
+        else:
+            token_ids_name, padding_mask_name = "token_ids", "padding_mask"
+
+        token_ids_input = keras.Input(
+            shape=(None,), dtype="int32", name=token_ids_name
+        )
+        padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name=padding_mask_name
+        )
+        inputs = {
+            token_ids_name: token_ids_input,
+            padding_mask_name: padding_mask_input,
+        }
+
+        if multimodal:
+            raw_image_size = self.vision_encoder.image_size
+            if isinstance(raw_image_size, (tuple, list)):
+                image_h, image_w = (
+                    int(raw_image_size[0]),
+                    int(raw_image_size[1]),
+                )
+            else:
+                image_h = image_w = int(raw_image_size)
+
+            images_input = keras.Input(
+                shape=(image_h, image_w, 3),
+                dtype="float32",
+                name="images",
+            )
+            inputs["images"] = images_input
+
+            patch_features = self.vision_encoder(images_input)
+            query_embeddings = self.qformer(patch_features)
+        else:
+            query_embeddings = None
+
+        # The language model itself always consumes `token_ids` /
+        # `padding_mask` for its (encoder) input; the backbone maps the
+        # externally exposed names onto these.
+        lm_inputs = {
+            "token_ids": token_ids_input,
+            "padding_mask": padding_mask_input,
+        }
+
+        if is_encoder_decoder:
+            decoder_token_ids_input = keras.Input(
+                shape=(None,), dtype="int32", name="decoder_token_ids"
+            )
+            decoder_padding_mask_input = keras.Input(
+                shape=(None,), dtype="int32", name="decoder_padding_mask"
+            )
+            inputs["decoder_token_ids"] = decoder_token_ids_input
+            inputs["decoder_padding_mask"] = decoder_padding_mask_input
+            lm_inputs["decoder_token_ids"] = decoder_token_ids_input
+            lm_inputs["decoder_padding_mask"] = decoder_padding_mask_input
+
+        if query_embeddings is not None:
+            lm_inputs["qformer_features"] = query_embeddings
+
+        output = self.language_model(lm_inputs)
+
+        super().__init__(
+            inputs=inputs,
+            outputs=output,
+            dtype=dtype,
+            **kwargs,
+        )
+
+        self.multimodal = multimodal
+
+    @property
+    def token_embedding(self):
+        """The token embedding layer of the language model."""
+        return self.language_model.token_embedding
+
+    @property
+    def num_query_tokens(self):
+        """Number of Q-Former query tokens (0 in text-only mode)."""
+        return self.qformer.num_query_tokens if self.qformer is not None else 0
+
+    @property
+    def qformer_hidden_dim(self):
+        """Hidden dimensionality of the Q-Former (0 in text-only mode)."""
+        return self.qformer.hidden_dim if self.qformer is not None else 0
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vision_encoder": (
+                    keras.layers.serialize(self.vision_encoder)
+                    if self.vision_encoder is not None
+                    else None
+                ),
+                "qformer": (
+                    keras.layers.serialize(self.qformer)
+                    if self.qformer is not None
+                    else None
+                ),
+                "language_model": keras.layers.serialize(self.language_model),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        config = dict(config)
+        for key in ("vision_encoder", "qformer"):
+            if config.get(key) is not None:
+                config[key] = keras.layers.deserialize(config[key])
+        config["language_model"] = keras.layers.deserialize(
+            config["language_model"]
+        )
+        return cls(**config)

keras_hub/src/models/blip2/blip2_backbone_test.py

@@ -0,0 +1,125 @@
+import numpy as np
+import pytest
+
+from keras_hub.src.models.blip2.blip2_backbone import BLIP2Backbone
+from keras_hub.src.models.blip2.blip2_custom_opt import BLIP2CustomOPT
+from keras_hub.src.models.blip2.blip2_qformer import BLIP2QFormer
+from keras_hub.src.models.blip2.blip2_vision_encoder import BLIP2VisionEncoder
+from keras_hub.src.tests.test_case import TestCase
+
+
+class BLIP2BackboneTest(TestCase):
+    def setUp(self):
+        self.batch_size = 2
+        self.image_size = 32
+        self.seq_length = 5
+        self.num_query_tokens = 2
+        self.hidden_dim = 4
+
+        vision_encoder = BLIP2VisionEncoder(
+            image_size=self.image_size,
+            patch_size=4,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=8,
+            intermediate_dim=16,
+            use_patch_bias=True,
+            use_class_token=True,
+            use_mha_bias=True,
+            use_mlp_bias=True,
+            dropout_rate=0.0,
+            layer_norm_epsilon=1e-6,
+            initializer_range=0.02,
+            dtype="float32",
+        )
+        qformer = BLIP2QFormer(
+            num_query_tokens=self.num_query_tokens,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=self.hidden_dim,
+            intermediate_dim=8,
+            vision_dim=8,
+            cross_attention_frequency=1,
+            dropout=0.0,
+            layer_norm_epsilon=1e-6,
+            dtype="float32",
+        )
+        language_model = BLIP2CustomOPT(
+            vocabulary_size=14,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=self.hidden_dim,
+            intermediate_dim=8,
+            num_query_tokens=self.num_query_tokens,
+            qformer_hidden_dim=self.hidden_dim,
+            max_sequence_length=10,
+            dropout=0.0,
+            language_projection=None,
+            dtype="float32",
+        )
+
+        self.init_kwargs = {
+            "vision_encoder": vision_encoder,
+            "qformer": qformer,
+            "language_model": language_model,
+        }
+        self.input_data = {
+            "images": np.ones(
+                (self.batch_size, self.image_size, self.image_size, 3),
+                dtype="float32",
+            ),
+            "token_ids": np.ones(
+                (self.batch_size, self.seq_length), dtype="int32"
+            ),
+            "padding_mask": np.ones(
+                (self.batch_size, self.seq_length), dtype="bool"
+            ),
+        }
+
+    def test_backbone_basics(self):
+        self.run_backbone_test(
+            cls=BLIP2Backbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+            expected_output_shape=(
+                self.batch_size,
+                self.num_query_tokens + self.seq_length,
+                self.hidden_dim,
+            ),
+            variable_length_data=[self.input_data],
+            run_quantization_check=False,
+            run_mixed_precision_check=False,
+        )
+
+    def test_architecture_characteristics(self):
+        backbone = BLIP2Backbone(**self.init_kwargs)
+        self.assertEqual(backbone.num_query_tokens, self.num_query_tokens)
+        self.assertEqual(backbone.qformer.hidden_dim, self.hidden_dim)
+        self.assertEqual(backbone.language_model.hidden_dim, self.hidden_dim)
+
+    @pytest.mark.large
+    def test_saved_model(self):
+        self.run_model_saving_test(
+            cls=BLIP2Backbone,
+            init_kwargs=self.init_kwargs,
+            input_data=self.input_data,
+        )
+
+    @pytest.mark.kaggle_key_required
+    @pytest.mark.extra_large
+    def test_smallest_preset(self):
+        self.run_preset_test(
+            cls=BLIP2Backbone,
+            preset="blip2_base",
+            input_data=self.input_data,
+        )
+
+    @pytest.mark.kaggle_key_required
+    @pytest.mark.extra_large
+    def test_all_presets(self):
+        for preset in BLIP2Backbone.presets:
+            self.run_preset_test(
+                cls=BLIP2Backbone,
+                preset=preset,
+                input_data=self.input_data,
+            )