keras-hub-nightly 0.16.1.dev202409250340__py3-none-any.whl → 0.16.1.dev202409260340__py3-none-any.whl

keras_hub/api/layers/__init__.py

@@ -50,6 +50,9 @@ from keras_hub.src.layers.preprocessing.resizing_image_converter import (
     ResizingImageConverter,
 )
 from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.densenet.densenet_image_converter import (
+    DenseNetImageConverter,
+)
 from keras_hub.src.models.pali_gemma.pali_gemma_image_converter import (
     PaliGemmaImageConverter,
 )

keras_hub/api/models/__init__.py

@@ -66,6 +66,8 @@ from keras_hub.src.models.bloom.bloom_causal_lm_preprocessor import (
 from keras_hub.src.models.bloom.bloom_tokenizer import BloomTokenizer
 from keras_hub.src.models.causal_lm import CausalLM
 from keras_hub.src.models.causal_lm_preprocessor import CausalLMPreprocessor
+from keras_hub.src.models.clip.clip_preprocessor import CLIPPreprocessor
+from keras_hub.src.models.clip.clip_tokenizer import CLIPTokenizer
 from keras_hub.src.models.csp_darknet.csp_darknet_backbone import (
     CSPDarkNetBackbone,
 )
@@ -100,6 +102,9 @@ from keras_hub.src.models.densenet.densenet_backbone import DenseNetBackbone
 from keras_hub.src.models.densenet.densenet_image_classifier import (
     DenseNetImageClassifier,
 )
+from keras_hub.src.models.densenet.densenet_image_classifier_preprocessor import (
+    DenseNetImageClassifierPreprocessor,
+)
 from keras_hub.src.models.distil_bert.distil_bert_backbone import (
     DistilBertBackbone,
 )
@@ -260,7 +265,17 @@ from keras_hub.src.models.sam.sam_backbone import SAMBackbone
 from keras_hub.src.models.sam.sam_image_segmenter import SAMImageSegmenter
 from keras_hub.src.models.seq_2_seq_lm import Seq2SeqLM
 from keras_hub.src.models.seq_2_seq_lm_preprocessor import Seq2SeqLMPreprocessor
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (
+    StableDiffusion3Backbone,
+)
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image import (
+    StableDiffusion3TextToImage,
+)
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (
+    StableDiffusion3TextToImagePreprocessor,
+)
 from keras_hub.src.models.t5.t5_backbone import T5Backbone
+from keras_hub.src.models.t5.t5_preprocessor import T5Preprocessor
 from keras_hub.src.models.t5.t5_tokenizer import T5Tokenizer
 from keras_hub.src.models.task import Task
 from keras_hub.src.models.text_classifier import TextClassifier
@@ -268,6 +283,7 @@ from keras_hub.src.models.text_classifier import TextClassifier as Classifier
 from keras_hub.src.models.text_classifier_preprocessor import (
     TextClassifierPreprocessor,
 )
+from keras_hub.src.models.text_to_image import TextToImage
 from keras_hub.src.models.vgg.vgg_backbone import VGGBackbone
 from keras_hub.src.models.vgg.vgg_image_classifier import VGGImageClassifier
 from keras_hub.src.models.vit_det.vit_det_backbone import ViTDetBackbone

keras_hub/api/tokenizers/__init__.py

@@ -21,6 +21,7 @@ from keras_hub.src.models.albert.albert_tokenizer import AlbertTokenizer
 from keras_hub.src.models.bart.bart_tokenizer import BartTokenizer
 from keras_hub.src.models.bert.bert_tokenizer import BertTokenizer
 from keras_hub.src.models.bloom.bloom_tokenizer import BloomTokenizer
+from keras_hub.src.models.clip.clip_tokenizer import CLIPTokenizer
 from keras_hub.src.models.deberta_v3.deberta_v3_tokenizer import (
     DebertaV3Tokenizer,
 )

keras_hub/src/models/{stable_diffusion_v3 → clip}/clip_encoder_block.py

@@ -11,6 +11,7 @@
 # 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.
+from keras import dtype_policies
 from keras import layers
 from keras import ops
 
@@ -43,7 +44,7 @@ class CLIPEncoderBlock(layers.Layer):
             intermediate_activation = quick_gelu
 
         self.layer_norm_1 = layers.LayerNormalization(
-            epsilon=0.00001, dtype=self.dtype_policy, name="layer_norm_1"
+            epsilon=1e-5, dtype="float32", name="layer_norm_1"
         )
         self.attention = layers.MultiHeadAttention(
             num_heads,
@@ -52,7 +53,7 @@ class CLIPEncoderBlock(layers.Layer):
             name="attention",
         )
         self.layer_norm_2 = layers.LayerNormalization(
-            epsilon=0.00001, dtype=self.dtype_policy, name="layer_norm_2"
+            epsilon=1e-5, dtype="float32", name="layer_norm_2"
         )
         self.dense_1 = layers.Dense(
             self.intermediate_dim, dtype=self.dtype_policy, name="dense_1"
@@ -67,6 +68,11 @@ class CLIPEncoderBlock(layers.Layer):
     def build(self, input_shape):
         self.layer_norm_1.build(input_shape)
         self.attention.build(input_shape, input_shape, input_shape)
+        # Before Keras 3.2, there was no setter for `dtype_policy`. Directly
+        # assign a `DTypePolicy` instead.
+        self.attention._softmax.dtype_policy = dtype_policies.DTypePolicy(
+            "float32"
+        )
         self.layer_norm_2.build(input_shape)
         self.dense_1.build(input_shape)
         input_shape = self.dense_1.compute_output_shape(input_shape)

keras_hub/src/models/clip/clip_preprocessor.py

@@ -0,0 +1,147 @@
+# Copyright 2024 The KerasHub Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.clip.clip_tokenizer import CLIPTokenizer
+from keras_hub.src.models.preprocessor import Preprocessor
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+try:
+    import tensorflow as tf
+except ImportError:
+    tf = None
+
+
+@keras_hub_export("keras_hub.models.CLIPPreprocessor")
+class CLIPPreprocessor(Preprocessor):
+    """CLIP preprocessing layer which tokenizes and packs inputs.
+
+    This preprocessing layer will do 2 things:
+
+    - Tokenize the inputs using the `tokenizer`.
+    - Construct a dictionary with keys `"token_ids"`, `"padding_mask"`.
+
+    This layer can be used directly with `tf.data.Dataset.map` to preprocess
+    string data in the `(x, y, sample_weight)` format used by
+    `keras.Model.fit`.
+
+    The call method of this layer accepts three arguments, `x`, `y`, and
+    `sample_weight`. `x` can be a python string or tensor representing a single
+    segment, a list of python strings representing a batch of single segments,
+    or a list of tensors representing multiple segments to be packed together.
+    `y` and `sample_weight` are both optional, can have any format, and will be
+    passed through unaltered.
+
+    `CLIPPreprocessor` forces the input to have only one segment, as CLIP is
+    mainly used for generation tasks. For tasks having multi-segment inputs
+    like "glue/mnli", please use a model designed for classification purposes
+    such as BERT or RoBERTa.
+
+    Args:
+        tokenizer: A `keras_hub.models.CLIPTokenizer` instance.
+        sequence_length: The length of the packed inputs.
+        add_start_token: If `True`, the preprocessor will prepend the tokenizer
+            start token to each input sequence.
+        add_end_token: If `True`, the preprocessor will append the tokenizer
+            end token to each input sequence.
+        to_lower: bool. Whether to lower the inputs.
+
+    Call arguments:
+        x: A string, `tf.Tensor` or list of python strings.
+        y: Any label data. Will be passed through unaltered.
+        sample_weight: Any label weight data. Will be passed through unaltered.
+        sequence_length: Pass to override the configured `sequence_length` of
+            the layer.
+    """
+
+    # TODO: Add example once we have a CLIP model.
+
+    tokenizer_cls = CLIPTokenizer
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length=77,
+        add_start_token=True,
+        add_end_token=True,
+        to_lower=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.tokenizer = tokenizer
+        self.packer = None
+        self.sequence_length = sequence_length
+        self.add_start_token = add_start_token
+        self.add_end_token = add_end_token
+        self.to_lower = to_lower
+
+    def build(self, input_shape):
+        # Defer packer creation to `build()` so that we can be sure tokenizer
+        # assets have loaded when restoring a saved model.
+        self.packer = StartEndPacker(
+            start_value=self.tokenizer.start_token_id,
+            end_value=self.tokenizer.end_token_id,
+            pad_value=self.tokenizer.end_token_id,
+            sequence_length=self.sequence_length,
+            return_padding_mask=True,
+        )
+        self.built = True
+
+    @preprocessing_function
+    def call(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        sequence_length=None,
+    ):
+        sequence_length = sequence_length or self.sequence_length
+        if self.to_lower:
+            x = tf.strings.lower(x)
+        token_ids, padding_mask = self.packer(
+            self.tokenizer(x),
+            sequence_length=sequence_length,
+            add_start_value=self.add_start_token,
+            add_end_value=self.add_end_token,
+        )
+        x = {
+            "token_ids": token_ids,
+            "padding_mask": padding_mask,
+        }
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "add_start_token": self.add_start_token,
+                "add_end_token": self.add_end_token,
+                "to_lower": self.to_lower,
+            }
+        )
+        return config
+
+    @property
+    def sequence_length(self):
+        """The padded length of model input sequences."""
+        return self._sequence_length
+
+    @sequence_length.setter
+    def sequence_length(self, value):
+        self._sequence_length = value
+        if self.packer is not None:
+            self.packer.sequence_length = value

keras_hub/src/models/{stable_diffusion_v3 → clip}/clip_text_encoder.py

@@ -11,21 +11,46 @@
 # 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_hub.src.layers.modeling.token_and_position_embedding import (
     TokenAndPositionEmbedding,
 )
-from keras_hub.src.models.stable_diffusion_v3.clip_encoder_block import (
-    CLIPEncoderBlock,
-)
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.clip.clip_encoder_block import CLIPEncoderBlock
+
 
+class CLIPTextEncoder(Backbone):
+    """CLIP text core network with hyperparameters.
+
+    Args:
+        vocabulary_size: int. The size of the token vocabulary.
+        embedding_dim: int. The output dimension of the embedding layer.
+        hidden_dim: int. The size of the transformer hidden state at the end
+            of each transformer layer.
+        num_layers: int. The number of transformer layers.
+        num_heads: int. The number of attention heads for each transformer.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        intermediate_dim: int. The output dimension of the first Dense layer in
+            a two-layer feedforward network for each transformer.
+        intermediate_activation: activation function. The activation that
+            is used for the first Dense layer in a two-layer feedforward network
+            for each transformer.
+        intermediate_output_index: optional int. The index of the intermediate
+            output. If specified, the output will become a dictionary with two
+            keys `"sequence_output"` and `"intermediate_output"`.
+        max_sequence_length: int. The maximum sequence length that this encoder
+            can consume.
+        dtype: string or `keras.mixed_precision.DTypePolicy`. The dtype to use
+            for the models computations and weights. Note that some
+            computations, such as softmax and layer normalization will always
+            be done a float32 precision regardless of dtype.
+    """
 
-class CLIPTextEncoder(keras.Model):
     def __init__(
         self,
+        vocabulary_size,
         embedding_dim,
         hidden_dim,
         num_layers,
@@ -33,9 +58,9 @@ class CLIPTextEncoder(keras.Model):
         intermediate_dim,
         intermediate_activation="quick_gelu",
         intermediate_output_index=None,
-        vocabulary_size=49408,
-        sequence_length=77,
+        max_sequence_length=77,
         dtype=None,
+        name=None,
         **kwargs,
     ):
         if (
@@ -44,13 +69,17 @@ class CLIPTextEncoder(keras.Model):
         ):
             intermediate_output_index += num_layers
 
+        # `prefix` is used to prevent duplicate name when utilizing multiple
+        # CLIP models within a single model, such as in StableDiffusion3.
+        prefix = str(name) + "_" if name is not None else ""
+
         # === Layers ===
         self.embedding = TokenAndPositionEmbedding(
             vocabulary_size=vocabulary_size,
-            sequence_length=sequence_length,
+            sequence_length=max_sequence_length,
             embedding_dim=embedding_dim,
             dtype=dtype,
-            name="embedding",
+            name=f"{prefix}embedding",
         )
         self.encoder_layers = [
             CLIPEncoderBlock(
@@ -59,59 +88,44 @@ class CLIPTextEncoder(keras.Model):
                 intermediate_dim,
                 intermediate_activation,
                 dtype=dtype,
+                name=f"{prefix}encoder_block_{i}",
             )
-            for _ in range(num_layers)
+            for i in range(num_layers)
         ]
         self.layer_norm = layers.LayerNormalization(
-            epsilon=0.00001, dtype=dtype, name="layer_norm"
-        )
-        self.text_projection = layers.Dense(
-            hidden_dim,
-            use_bias=False,
-            dtype=dtype,
-            name="text_projection",
+            epsilon=1e-6, dtype="float32", name=f"{prefix}layer_norm"
         )
 
         # === Functional Model ===
-        encoder_token_ids = layers.Input(
-            shape=(sequence_length,), dtype="int32", name="encoder_token_ids"
+        token_id_input = layers.Input(
+            shape=(None,), dtype="int32", name="token_ids"
         )
-        x = self.embedding(encoder_token_ids)
-        encoder_intermediate_output = None
-        # Encoder.
+        x = self.embedding(token_id_input)
+        intermediate_output = None
         for i, block in enumerate(self.encoder_layers):
             x = block(x)
             if i == intermediate_output_index:
-                encoder_intermediate_output = x
+                intermediate_output = x
         x = self.layer_norm(x)
-        encoder_output = x
-        if encoder_intermediate_output is not None:
-            encoder_intermediate_output = self.layer_norm(
-                encoder_intermediate_output
-            )
-        # Projection.
-        indices = ops.expand_dims(
-            ops.cast(ops.argmax(encoder_token_ids, axis=-1), "int32"), axis=-1
-        )
-        pooled_output = ops.take_along_axis(x, indices[:, :, None], axis=1)
-        pooled_output = ops.squeeze(pooled_output, axis=1)
-        projection_output = self.text_projection(pooled_output)
+        sequence_output = x
 
-        outputs = {
-            "encoder_sequence_output": encoder_output,
-            "encoder_pooled_output": pooled_output,
-            "encoder_projection_output": projection_output,
-        }
         if intermediate_output_index is not None:
-            outputs["encoder_intermediate_output"] = encoder_intermediate_output
-
+            outputs = {
+                "sequence_output": sequence_output,
+                "intermediate_output": intermediate_output,
+            }
+        else:
+            outputs = sequence_output
         super().__init__(
-            inputs={"encoder_token_ids": encoder_token_ids},
+            inputs={"token_ids": token_id_input},
             outputs=outputs,
+            name=name,
             **kwargs,
         )
 
         # === Config ===
+        self.vocabulary_size = vocabulary_size
+        self.max_sequence_length = max_sequence_length
         self.embedding_dim = embedding_dim
         self.hidden_dim = hidden_dim
         self.num_layers = num_layers
@@ -119,22 +133,12 @@ class CLIPTextEncoder(keras.Model):
         self.intermediate_dim = intermediate_dim
         self.intermediate_activation = intermediate_activation
         self.intermediate_output_index = intermediate_output_index
-        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(
             {
+                "vocabulary_size": self.vocabulary_size,
                 "embedding_dim": self.embedding_dim,
                 "hidden_dim": self.hidden_dim,
                 "num_layers": self.num_layers,
@@ -142,8 +146,7 @@ class CLIPTextEncoder(keras.Model):
                 "intermediate_dim": self.intermediate_dim,
                 "intermediate_activation": self.intermediate_activation,
                 "intermediate_output_index": self.intermediate_output_index,
-                "vocabulary_size": self.vocabulary_size,
-                "sequence_length": self.sequence_length,
+                "max_sequence_length": self.max_sequence_length,
             }
         )
         return config

keras_hub/src/models/{stable_diffusion_v3 → clip}/clip_tokenizer.py

@@ -11,9 +11,12 @@
 # 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.
+
+from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.tokenizers.byte_pair_tokenizer import BytePairTokenizer
 from keras_hub.src.tokenizers.byte_pair_tokenizer import convert_to_ragged_batch
 from keras_hub.src.tokenizers.byte_pair_tokenizer import split_strings_for_bpe
+from keras_hub.src.utils.tensor_utils import preprocessing_function
 
 try:
     import tensorflow as tf
@@ -21,10 +24,51 @@ except ImportError:
     tf = None
 
 
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.CLIPTokenizer",
+        "keras_hub.models.CLIPTokenizer",
+    ]
+)
 class CLIPTokenizer(BytePairTokenizer):
-    def __init__(self, vocabulary=None, merges=None, **kwargs):
-        self.start_token = "<|startoftext|>"
-        self.end_token = "<|endoftext|>"
+    """A CLIP tokenizer using Byte-Pair Encoding subword segmentation.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_hub.tokenizers.BytePairTokenizer`. Unlike the
+    underlying tokenizer, it will check for all special tokens needed by CLIP
+    models and provides a `from_preset()` method to automatically download
+    a matching vocabulary for a CLIP preset.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        vocabulary: string or dict, maps token to integer ids. If it is a
+            string, it should be the file path to a json file.
+        merges: string or list, contains the merge rule. If it is a string,
+            it should be the file path to merge rules. The merge rule file
+            should have one merge rule per line. Every merge rule contains
+            merge entities separated by a space.
+        pad_with_end_token: bool. Whether to pad the output with `end_token`.
+    """
+
+    # TODO: Add example and `backbone_cls` once we have a CLIP model.
+
+    backbone_cls = None
+
+    def __init__(
+        self,
+        vocabulary=None,
+        merges=None,
+        pad_with_end_token=False,
+        **kwargs,
+    ):
+        self._add_special_token("<|startoftext|>", "start_token")
+        self._add_special_token("<|endoftext|>", "end_token")
+        self.pad_token_id = 0
 
         super().__init__(
             vocabulary=vocabulary,
@@ -33,35 +77,21 @@ class CLIPTokenizer(BytePairTokenizer):
             **kwargs,
         )
 
-    def set_vocabulary_and_merges(self, vocabulary, merges):
-        super().set_vocabulary_and_merges(vocabulary, merges)
-
-        if vocabulary is not None:
-            # Check for necessary special tokens.
-            if self.end_token not in self.get_vocabulary():
-                raise ValueError(
-                    f"Cannot find token `'{self.end_token}'` in the provided "
-                    f"`vocabulary`. Please provide `'{self.end_token}'` in "
-                    "your `vocabulary` or use a pretrained `vocabulary` name."
-                )
-
-            self.start_token_id = self.token_to_id(self.start_token)
-            self.end_token_id = self.token_to_id(self.end_token)
-            self.pad_token_id = 0
-        else:
-            self.end_token_id = None
-            self.start_token_id = None
-            self.pad_token_id = None
+        # When `pad_with_end_token` is True, we need to access the vocabulary,
+        # so the check is required.
+        if pad_with_end_token:
+            self._check_vocabulary()
+            self.pad_token_id = self.end_token_id
+        self.pad_with_end_token = pad_with_end_token
 
     def _bpe_merge_and_update_cache(self, tokens):
         """Process unseen tokens and add to cache."""
         words = self._transform_bytes(tokens)
 
-        # In StableDiffusionV3, we need to add `</w>` to the last word.
+        # In CLIP, we need to add `</w>` to the last word.
         words = tf.strings.reduce_join(words, axis=1, separator=" ")
         words = tf.strings.join([words, "</w>"])
         words = tf.strings.split(words, sep=" ")
-
         tokenized_words = self._bpe_merge(words)
 
         # For each word, join all its token by a whitespace,
@@ -71,17 +101,20 @@ class CLIPTokenizer(BytePairTokenizer):
         )
         self.cache.insert(tokens, tokenized_words)
 
+    @preprocessing_function
     def tokenize(self, inputs):
         self._check_vocabulary()
-        if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
-            inputs = tf.convert_to_tensor(inputs)
-
         if self.add_prefix_space:
             inputs = tf.strings.join([" ", inputs])
 
-        scalar_input = inputs.shape.rank == 0
-        if scalar_input:
+        unbatched = inputs.shape.rank == 0
+        if unbatched:
             inputs = tf.expand_dims(inputs, 0)
+        if inputs.shape.rank > 1:
+            raise ValueError(
+                "`tokenize()` inputs should be a string, list of strings, or "
+                f"string tensor with rank < 2. Received: {inputs}"
+            )
 
         raw_tokens = split_strings_for_bpe(inputs, self.unsplittable_tokens)
 
@@ -131,12 +164,13 @@ class CLIPTokenizer(BytePairTokenizer):
             tokens = tokens.to_tensor(shape=output_shape)
 
         # Convert to a dense output if input in scalar
-        if scalar_input:
+        if unbatched:
             tokens = tf.squeeze(tokens, 0)
             tf.ensure_shape(tokens, shape=[self.sequence_length])
 
         return tokens
 
+    @preprocessing_function
     def detokenize(self, inputs):
         self._check_vocabulary()
         inputs, unbatched, _ = convert_to_ragged_batch(inputs)
@@ -160,6 +194,11 @@ class CLIPTokenizer(BytePairTokenizer):
 
     def get_config(self):
         config = super().get_config()
+        config.update(
+            {
+                "pad_with_end_token": self.pad_with_end_token,
+            }
+        )
         # In the constructor, we pass the list of special tokens to the
         # `unsplittable_tokens` arg of the superclass' constructor. Hence, we
         # delete it from the config here.

keras_hub/src/models/densenet/__init__.py

@@ -11,3 +11,9 @@
 # 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.
+
+from keras_hub.src.models.densenet.densenet_backbone import DenseNetBackbone
+from keras_hub.src.models.densenet.densenet_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, DenseNetBackbone)

keras_hub/src/models/densenet/densenet_backbone.py

@@ -92,11 +92,14 @@ class DenseNetBackbone(FeaturePyramidBackbone):
                 channel_axis,
                 stackwise_num_repeats[stack_index],
                 growth_rate,
-                name=f"conv{index}",
+                name=f"stack{stack_index+1}",
             )
             pyramid_outputs[f"P{index}"] = x
             x = apply_transition_block(
-                x, channel_axis, compression_ratio, name=f"pool{index}"
+                x,
+                channel_axis,
+                compression_ratio,
+                name=f"transition{stack_index+1}",
             )
 
         x = apply_dense_block(
@@ -104,7 +107,7 @@ class DenseNetBackbone(FeaturePyramidBackbone):
             channel_axis,
             stackwise_num_repeats[-1],
             growth_rate,
-            name=f"conv{len(stackwise_num_repeats) + 1}",
+            name=f"stack{len(stackwise_num_repeats)}",
         )
         pyramid_outputs[f"P{len(stackwise_num_repeats) + 1}"] = x
         x = keras.layers.BatchNormalization(
@@ -148,7 +151,7 @@ def apply_dense_block(x, channel_axis, num_repeats, growth_rate, name=None):
 
     for i in range(num_repeats):
         x = apply_conv_block(
-            x, channel_axis, growth_rate, name=f"{name}_block_{i}"
+            x, channel_axis, growth_rate, name=f"{name}_block{i+1}"
         )
     return x
 
@@ -196,9 +199,9 @@ def apply_conv_block(x, channel_axis, growth_rate, name=None):
 
     shortcut = x
     x = keras.layers.BatchNormalization(
-        axis=channel_axis, epsilon=BN_EPSILON, name=f"{name}_0_bn"
+        axis=channel_axis, epsilon=BN_EPSILON, name=f"{name}_1_bn"
     )(x)
-    x = keras.layers.Activation("relu", name=f"{name}_0_relu")(x)
+    x = keras.layers.Activation("relu", name=f"{name}_1_relu")(x)
     x = keras.layers.Conv2D(
         4 * growth_rate,
         1,
@@ -207,9 +210,9 @@ def apply_conv_block(x, channel_axis, growth_rate, name=None):
         name=f"{name}_1_conv",
     )(x)
     x = keras.layers.BatchNormalization(
-        axis=channel_axis, epsilon=BN_EPSILON, name=f"{name}_1_bn"
+        axis=channel_axis, epsilon=BN_EPSILON, name=f"{name}_2_bn"
     )(x)
-    x = keras.layers.Activation("relu", name=f"{name}_1_relu")(x)
+    x = keras.layers.Activation("relu", name=f"{name}_2_relu")(x)
     x = keras.layers.Conv2D(
         growth_rate,
         3,