Add include rescaling to the pali gemma backbone (#1650)

To disable this option, pass

```
keras_nlp.models.PaliGemma.from_preset(
    "pali_gemma_3b_224",
    include_rescaling=False,
)
```

Allow inputs to be the more standard range 0, 255

Author: mattdangerw

keras_nlp/src/models/pali_gemma/pali_gemma_backbone.py

@@ -74,6 +74,8 @@ class PaliGemmaBackbone(Backbone):
         vit_classifier_activation: activation function. The activation that
             is used for final output classification in the vision transformer.
         vit_name: string. The name used for vision transformer layers.
+        include_rescaling: bool. If true, the image input will be rescaled from
+            the range `[0, 255]`, to the range `[0, 1]`.
         layer_norm_epsilon: float. The epsilon value user for every layer norm
             in all transformer blocks.
         dropout: float. Dropout probability for the Transformer decoder blocks.
@@ -132,6 +134,7 @@ def __init__(
         vit_pooling=None,
         vit_classifier_activation=None,
         vit_name=None,
+        include_rescaling=True,
         layer_norm_epsilon=1e-6,
         dropout=0,
         dtype=None,
@@ -163,6 +166,7 @@ def __init__(
         vit_intermediate_dim = vit_intermediate_dim or 4304
         self.vit_encoder = PaliGemmaVit(
             image_size=image_size,
+            include_rescaling=include_rescaling,
             patch_size=vit_patch_size,
             num_heads=vit_num_heads,
             hidden_dim=vit_hidden_dim,
@@ -232,6 +236,7 @@ def __init__(
         # === Config ===
         self.vocabulary_size = vocabulary_size
         self.image_size = image_size
+        self.include_rescaling = include_rescaling
         self.num_layers = num_layers
         self.num_query_heads = num_query_heads
         self.num_key_value_heads = num_key_value_heads
@@ -258,6 +263,7 @@ def get_config(self):
             {
                 "vocabulary_size": self.vocabulary_size,
                 "image_size": self.image_size,
+                "include_rescaling": self.include_rescaling,
                 "num_layers": self.num_layers,
                 "num_query_heads": self.num_query_heads,
                 "num_key_value_heads": self.num_key_value_heads,

keras_nlp/src/models/pali_gemma/pali_gemma_vit.py

@@ -423,6 +423,8 @@ class PaliGemmaVit(keras.Model):
     Args:
         image_size: int. The height/width of the image. Both height and width is
             expected to be the same.
+        include_rescaling: bool. If true, the image input will be rescaled from
+            the range `[0, 255]`, to the range `[0, 1]`.
         patch_size: int. The size of each square patch in the input image.
         num_heads: int. The number of attention heads for the vision(image)
             transformer encoder.
@@ -463,6 +465,7 @@ def __init__(
         num_layers,
         intermediate_dim,
         num_classes,
+        include_rescaling=True,
         pooling=None,
         classifier_activation=None,
         dtype=None,
@@ -472,7 +475,13 @@ def __init__(
         image_input = keras.Input(
             shape=(image_size, image_size, 3), name="images"
         )
-        encoded = PaliGemmaVitEncoder(
+        x = image_input  # Intermediate result.
+        if include_rescaling:
+            rescaling = keras.layers.Rescaling(
+                scale=1.0 / 127.5, offset=-1.0, name="rescaling"
+            )
+            x = rescaling(image_input)
+        x = PaliGemmaVitEncoder(
             hidden_dim=hidden_dim,
             num_layers=num_layers,
             num_heads=num_heads,
@@ -481,20 +490,20 @@ def __init__(
             image_size=image_size,
             dtype=dtype,
             name="image_encoder",
-        )(image_input)
+        )(x)
         if pooling == "map":
-            pooled = MultiHeadAttentionPooling(
+            x = MultiHeadAttentionPooling(
                 num_heads=num_heads,
                 hidden_dim=hidden_dim,
                 dtype=dtype,
                 name="pooling",
-            )(encoded)
+            )(x)
         elif pooling == "gap":
-            pooled = ops.mean(encoded, axis=1)
+            x = ops.mean(x, axis=1)
         elif pooling == "zero":
-            pooled = encoded[:, 0]
+            x = x[:, 0]
         elif pooling is None:
-            pooled = encoded
+            x = x
         else:
             raise ValueError(
                 "Invalid value for argument `pooling`. "
@@ -506,7 +515,7 @@ def __init__(
             activation=classifier_activation,
             dtype=dtype,
             name="image_classifier",
-        )(pooled)
+        )(x)
         super().__init__(
             inputs=image_input,
             outputs=outputs,
@@ -521,6 +530,7 @@ def __init__(
         self.pooling = pooling
         self.num_classes = num_classes
         self.image_size = image_size
+        self.include_rescaling = include_rescaling
         self.patch_size = patch_size
         self.classifier_activation = keras.activations.get(
             classifier_activation
@@ -541,6 +551,7 @@ def get_config(self):
                     self.classifier_activation
                 ),
                 "image_size": self.image_size,
+                "include_rescaling": self.include_rescaling,
                 "patch_size": self.patch_size,
             }
         )

keras_nlp/src/models/pali_gemma/pali_gemma_vit_test.py

@@ -45,6 +45,23 @@ def test_vit_encoder(self):
             output.shape, (batch_size, intermediate_dim, hidden_dim)
         )
 
+    def test_vit_rescaling(self):
+        vit_encoder = PaliGemmaVit(
+            image_size=16,
+            patch_size=4,
+            hidden_dim=8,
+            num_layers=2,
+            num_heads=2,
+            intermediate_dim=16,
+            num_classes=32,
+        )
+        self.assertIsNotNone(vit_encoder.get_layer("rescaling"))
+        with self.assertRaises(ValueError):
+            config = vit_encoder.get_config()
+            config["include_rescaling"] = False
+            vit_encoder = PaliGemmaVit.from_config(config)
+            vit_encoder.get_layer("rescaling")
+
     def test_vision_embeddings(self):
         embeddings_layer = PaliGemmaVitEmbeddings(
             image_size=16,