Replace Backbone with keras.Model in CLIPTextEncoder and T5XXLTextEncoder (#1802)

Author: james77777778

keras_nlp/src/models/stable_diffusion_v3/clip_text_encoder.py

@@ -11,19 +11,19 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import keras
 from keras import layers
 from keras import ops
 
 from keras_nlp.src.layers.modeling.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
-from keras_nlp.src.models.backbone import Backbone
 from keras_nlp.src.models.stable_diffusion_v3.clip_encoder_block import (
     CLIPEncoderBlock,
 )
 
 
-class CLIPTextEncoder(Backbone):
+class CLIPTextEncoder(keras.Model):
     def __init__(
         self,
         embedding_dim,
@@ -108,7 +108,6 @@ def __init__(
         super().__init__(
             inputs={"encoder_token_ids": encoder_token_ids},
             outputs=outputs,
-            dtype=dtype,
             **kwargs,
         )
 
@@ -123,6 +122,15 @@ def __init__(
         self.vocabulary_size = vocabulary_size
         self.sequence_length = sequence_length
 
+        if dtype is not None:
+            try:
+                self.dtype_policy = keras.dtype_policies.get(dtype)
+            # Before Keras 3.2, there is no `keras.dtype_policies.get`.
+            except AttributeError:
+                if isinstance(dtype, keras.DTypePolicy):
+                    dtype = dtype.name
+                self.dtype_policy = keras.DTypePolicy(dtype)
+
     def get_config(self):
         config = super().get_config()
         config.update(

keras_nlp/src/models/stable_diffusion_v3/t5_xxl_text_encoder.py

@@ -16,12 +16,11 @@
 from keras_nlp.src.layers.modeling.reversible_embedding import (
     ReversibleEmbedding,
 )
-from keras_nlp.src.models.backbone import Backbone
 from keras_nlp.src.models.t5.t5_layer_norm import T5LayerNorm
 from keras_nlp.src.models.t5.t5_transformer_layer import T5TransformerLayer
 
 
-class T5XXLTextEncoder(Backbone):
+class T5XXLTextEncoder(keras.Model):
     def __init__(
         self,
         vocabulary_size,
@@ -111,7 +110,6 @@ def __init__(
                 "encoder_padding_mask": encoder_padding_mask_input,
             },
             outputs=encoder_output,
-            dtype=dtype,
             **kwargs,
         )
 
@@ -128,6 +126,15 @@ def __init__(
         self.layer_norm_epsilon = layer_norm_epsilon
         self.tie_embedding_weights = tie_embedding_weights
 
+        if dtype is not None:
+            try:
+                self.dtype_policy = keras.dtype_policies.get(dtype)
+            # Before Keras 3.2, there is no `keras.dtype_policies.get`.
+            except AttributeError:
+                if isinstance(dtype, keras.DTypePolicy):
+                    dtype = dtype.name
+                self.dtype_policy = keras.DTypePolicy(dtype)
+
     def get_config(self):
         config = super().get_config()
         config.update(