keras-hub-nightly 0.19.0.dev202412120352__py3-none-any.whl → 0.19.0.dev202412140350__py3-none-any.whl

keras_hub/src/models/sam/sam_transformer.py

@@ -31,14 +31,15 @@ class TwoWayTransformer(keras.layers.Layer):
     location and type.
 
     Args:
-        num_layers: int, optional. The num_layers of the attention blocks (the number
-            of attention blocks to use). Defaults to `2`.
+        num_layers: int, optional. The num_layers of the attention blocks
+            (the number of attention blocks to use). Defaults to `2`.
         hidden_size: int, optional. The number of features of the input image
             and point embeddings. Defaults to `256`.
         num_heads: int, optional. Number of heads to use in the attention
             layers. Defaults to `8`.
-        intermediate_dim: int, optional. The number of units in the hidden layer of
-            the MLP block used in the attention layers. Defaults to `2048`.
+        intermediate_dim: int, optional. The number of units in the hidden
+            layer of the MLP block used in the attention layers.
+            Defaults to `2048`.
         activation: str, optional. The activation of the MLP block's output
             layer used in the attention layers. Defaults to `"relu"`.
         attention_downsample_rate: int, optional. The downsample rate of the

keras_hub/src/models/segformer/segformer_backbone.py

@@ -6,18 +6,19 @@ from keras_hub.src.models.backbone import Backbone
 
 @keras_hub_export("keras_hub.models.SegFormerBackbone")
 class SegFormerBackbone(Backbone):
-    """A Keras model implementing the SegFormer architecture for semantic segmentation.
+    """A Keras model implementing SegFormer for semantic segmentation.
 
-    This class implements the majority of the SegFormer architecture described in
-    [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers]
-    (https://arxiv.org/abs/2105.15203) and [based on the TensorFlow implementation from DeepVision]
-    (https://github.com/DavidLandup0/deepvision/tree/main/deepvision/models/segmentation/segformer).
+    This class implements the majority of the SegFormer architecture described
+    in [SegFormer: Simple and Efficient Design for Semantic Segmentation](https://arxiv.org/abs/2105.15203)
+    and based on the TensorFlow implementation
+    [from DeepVision](https://github.com/DavidLandup0/deepvision/tree/main/deepvision/models/segmentation/segformer).
 
-    SegFormers are meant to be used with the MixTransformer (MiT) encoder family, and
-    and use a very lightweight all-MLP decoder head.
+    SegFormers are meant to be used with the MixTransformer (MiT) encoder
+    family, and use a very lightweight all-MLP decoder head.
 
-    The MiT encoder uses a hierarchical transformer which outputs features at multiple scales,
-    similar to that of the hierarchical outputs typically associated with CNNs.
+    The MiT encoder uses a hierarchical transformer which outputs features at
+    multiple scales, similar to that of the hierarchical outputs typically
+    associated with CNNs.
 
     Args:
         image_encoder: `keras.Model`. The backbone network for the model that is
@@ -50,7 +51,8 @@ class SegFormerBackbone(Backbone):
         strides=[4, 2, 2, 2],
     )
 
-    segformer_backbone = keras_hub.models.SegFormerBackbone(image_encoder=backbone, projection_filters=256)
+    segformer_backbone = keras_hub.models.SegFormerBackbone(
+        image_encoder=backbone, projection_filters=256)
     ```
 
     Using the class with a preset `backbone`:
@@ -59,7 +61,8 @@ class SegFormerBackbone(Backbone):
     import keras_hub
 
     backbone = keras_hub.models.MiTBackbone.from_preset("mit_b0_ade20k_512")
-    segformer_backbone = keras_hub.models.SegFormerBackbone(image_encoder=backbone, projection_filters=256)
+    segformer_backbone = keras_hub.models.SegFormerBackbone(
+        image_encoder=backbone, projection_filters=256)
     ```
 
     """
@@ -74,9 +77,10 @@ class SegFormerBackbone(Backbone):
             image_encoder, keras.Model
         ):
             raise ValueError(
-                "Argument `image_encoder` must be a `keras.layers.Layer` instance "
-                f" or `keras.Model`. Received instead "
-                f"image_encoder={image_encoder} (of type {type(image_encoder)})."
+                "Argument `image_encoder` must be a `keras.layers.Layer` "
+                f"instance or `keras.Model`. Received instead "
+                f"image_encoder={image_encoder} "
+                f"(of type {type(image_encoder)})."
             )
 
         # === Layers ===

keras_hub/src/models/segformer/segformer_image_segmenter.py

@@ -3,41 +3,43 @@ import keras
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.image_segmenter import ImageSegmenter
 from keras_hub.src.models.segformer.segformer_backbone import SegFormerBackbone
-from keras_hub.src.models.segformer.segformer_image_segmenter_preprocessor import (
+from keras_hub.src.models.segformer.segformer_image_segmenter_preprocessor import (  # noqa: E501
     SegFormerImageSegmenterPreprocessor,
 )
 
 
 @keras_hub_export("keras_hub.models.SegFormerImageSegmenter")
 class SegFormerImageSegmenter(ImageSegmenter):
-    """A Keras model implementing the SegFormer architecture for semantic segmentation.
+    """A Keras model implementing SegFormer for semantic segmentation.
 
-    This class implements the segmentation head of the SegFormer architecture described in
-    [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers]
-    (https://arxiv.org/abs/2105.15203) and [based on the TensorFlow implementation from DeepVision]
+    This class implements the segmentation head of the SegFormer architecture
+    described in [SegFormer: Simple and Efficient Design for Semantic
+    Segmentation with Transformers] (https://arxiv.org/abs/2105.15203) and
+    [based on the TensorFlow implementation from DeepVision]
     (https://github.com/DavidLandup0/deepvision/tree/main/deepvision/models/segmentation/segformer).
 
-    SegFormers are meant to be used with the MixTransformer (MiT) encoder family, and
-    and use a very lightweight all-MLP decoder head.
+    SegFormers are meant to be used with the MixTransformer (MiT) encoder
+    family, and and use a very lightweight all-MLP decoder head.
 
-    The MiT encoder uses a hierarchical transformer which outputs features at multiple scales,
-    similar to that of the hierarchical outputs typically associated with CNNs.
+    The MiT encoder uses a hierarchical transformer which outputs features at
+    multiple scales, similar to that of the hierarchical outputs typically
+    associated with CNNs.
 
     Args:
         image_encoder: `keras.Model`. The backbone network for the model that is
-            used as a feature extractor for the SegFormer encoder.
-            It is *intended* to be used only with the MiT backbone model
-            (`keras_hub.models.MiTBackbone`) which was created
-            specifically for SegFormers.
-            Alternatively, can be a `keras_hub.models.Backbone` a model subclassing
-            `keras_hub.models.FeaturePyramidBackbone`, or a `keras.Model`
-            that has a `pyramid_outputs` property which is
-            a dictionary with keys "P2", "P3", "P4", and "P5" and layer names as values.
+            used as a feature extractor for the SegFormer encoder. It is
+            *intended* to be used only with the MiT backbone model
+            (`keras_hub.models.MiTBackbone`) which was created specifically for
+            SegFormers. Alternatively, can be a `keras_hub.models.Backbone` a
+            model subclassing `keras_hub.models.FeaturePyramidBackbone`, or a
+            `keras.Model` that has a `pyramid_outputs` property which is a
+            dictionary with keys "P2", "P3", "P4", and "P5" and layer names as
+            values.
         num_classes: int, the number of classes for the detection model,
             including the background class.
         projection_filters: int, number of filters in the
-            convolution layer projecting the concatenated features into
-            a segmentation map. Defaults to 256`.
+            convolution layer projecting the concatenated features into a
+            segmentation map. Defaults to 256`.
 
 
     Example:
@@ -45,10 +47,9 @@ class SegFormerImageSegmenter(ImageSegmenter):
     Using presets:
 
     ```python
-    import keras_hub
-    import numpy as np
-
-    segmenter = keras_hub.models.SegFormerImageSegmenter.from_preset("segformer_b0_ade20k_512")
+    segmenter = keras_hub.models.SegFormerImageSegmenter.from_preset(
+        "segformer_b0_ade20k_512"
+    )
 
     images = np.random.rand(1, 512, 512, 3)
     segformer(images)
@@ -57,17 +58,18 @@ class SegFormerImageSegmenter(ImageSegmenter):
     Using the SegFormer backbone:
 
     ```python
-    encoder = keras_hub.models.MiTBackbone.from_preset("mit_b0_ade20k_512")
-    backbone = keras_hub.models.SegFormerBackbone(image_encoder=encoder, projection_filters=256)
+    encoder = keras_hub.models.MiTBackbone.from_preset(
+        "mit_b0_ade20k_512"
+    )
+    backbone = keras_hub.models.SegFormerBackbone(
+        image_encoder=encoder,
+        projection_filters=256,
+    )
     ```
 
     Using the SegFormer backbone with a custom encoder:
 
     ```python
-    import keras
-    import keras_hub
-    import numpy as np
-
     images = np.ones(shape=(1, 96, 96, 3))
     labels = np.zeros(shape=(1, 96, 96, 1))
 
@@ -83,20 +85,31 @@ class SegFormerImageSegmenter(ImageSegmenter):
         strides=[4, 2, 2, 2],
     )
 
-    backbone = keras_hub.models.SegFormerBackbone(image_encoder=encoder, projection_filters=256)
-    segformer = keras_hub.models.SegFormerImageSegmenter(backbone=backbone, num_classes=4)
-
-    segformer(images)
+    backbone = keras_hub.models.SegFormerBackbone(
+        image_encoder=encoder,
+        projection_filters=256,
+    )
+    segformer = keras_hub.models.SegFormerImageSegmenter(
+        backbone=backbone,
+        num_classes=4,
+    )
+    segformer(images
     ```
 
     Using the segmentor class with a preset backbone:
 
     ```python
-    import keras_hub
-
-    image_encoder = keras_hub.models.MiTBackbone.from_preset("mit_b0_ade20k_512")
-    backbone = keras_hub.models.SegFormerBackbone(image_encoder=encoder, projection_filters=256)
-    segformer = keras_hub.models.SegFormerImageSegmenter(backbone=backbone, num_classes=4)
+    image_encoder = keras_hub.models.MiTBackbone.from_preset(
+        "mit_b0_ade20k_512"
+    )
+    backbone = keras_hub.models.SegFormerBackbone(
+        image_encoder=encoder,
+        projection_filters=256,
+    )
+    segformer = keras_hub.models.SegFormerImageSegmenter(
+        backbone=backbone,
+        num_classes=4,
+    )
     ```
     """
 

keras_hub/src/models/segformer/segformer_presets.py

@@ -4,7 +4,8 @@ presets = {
     "segformer_b0_ade20k_512": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB0 backbone fine-tuned on ADE20k in 512x512 resolution."
+                "SegFormer model with MiTB0 backbone fine-tuned on ADE20k in "
+                "512x512 resolution."
             ),
             "params": 3719027,
             "path": "segformer_b0",
@@ -14,7 +15,8 @@ presets = {
     "segformer_b1_ade20k_512": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB1 backbone fine-tuned on ADE20k in 512x512 resolution."
+                "SegFormer model with MiTB1 backbone fine-tuned on ADE20k in "
+                "512x512 resolution."
             ),
             "params": 13682643,
             "path": "segformer_b1",
@@ -24,7 +26,8 @@ presets = {
     "segformer_b2_ade20k_512": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB2 backbone fine-tuned on ADE20k in 512x512 resolution."
+                "SegFormer model with MiTB2 backbone fine-tuned on ADE20k in "
+                "512x512 resolution."
             ),
             "params": 24727507,
             "path": "segformer_b2",
@@ -34,7 +37,8 @@ presets = {
     "segformer_b3_ade20k_512": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB3 backbone fine-tuned on ADE20k in 512x512 resolution."
+                "SegFormer model with MiTB3 backbone fine-tuned on ADE20k in "
+                "512x512 resolution."
             ),
             "params": 44603347,
             "path": "segformer_b3",
@@ -44,7 +48,8 @@ presets = {
     "segformer_b4_ade20k_512": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB4 backbone fine-tuned on ADE20k in 512x512 resolution."
+                "SegFormer model with MiTB4 backbone fine-tuned on ADE20k in "
+                "512x512 resolution."
             ),
             "params": 61373907,
             "path": "segformer_b4",
@@ -54,7 +59,8 @@ presets = {
     "segformer_b5_ade20k_640": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB5 backbone fine-tuned on ADE20k in 640x640 resolution."
+                "SegFormer model with MiTB5 backbone fine-tuned on ADE20k in "
+                "640x640 resolution."
             ),
             "params": 81974227,
             "path": "segformer_b5",
@@ -64,7 +70,8 @@ presets = {
     "segformer_b0_cityscapes_1024": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB0 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+                "SegFormer model with MiTB0 backbone fine-tuned on Cityscapes "
+                "in 1024x1024 resolution."
             ),
             "params": 3719027,
             "path": "segformer_b0",
@@ -74,7 +81,8 @@ presets = {
     "segformer_b1_cityscapes_1024": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB1 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+                "SegFormer model with MiTB1 backbone fine-tuned on Cityscapes "
+                "in 1024x1024 resolution."
             ),
             "params": 13682643,
             "path": "segformer_b1",
@@ -84,7 +92,8 @@ presets = {
     "segformer_b2_cityscapes_1024": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB2 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+                "SegFormer model with MiTB2 backbone fine-tuned on Cityscapes "
+                "in 1024x1024 resolution."
             ),
             "params": 24727507,
             "path": "segformer_b2",
@@ -94,7 +103,8 @@ presets = {
     "segformer_b3_cityscapes_1024": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB3 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+                "SegFormer model with MiTB3 backbone fine-tuned on Cityscapes "
+                "in 1024x1024 resolution."
             ),
             "params": 44603347,
             "path": "segformer_b3",
@@ -104,7 +114,8 @@ presets = {
     "segformer_b4_cityscapes_1024": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB4 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+                "SegFormer model with MiTB4 backbone fine-tuned on Cityscapes "
+                "in 1024x1024 resolution."
             ),
             "params": 61373907,
             "path": "segformer_b4",
@@ -114,7 +125,8 @@ presets = {
     "segformer_b5_cityscapes_1024": {
         "metadata": {
             "description": (
-                "SegFormer model with MiTB5 backbone fine-tuned on Cityscapes in 1024x1024 resolution."
+                "SegFormer model with MiTB5 backbone fine-tuned on Cityscapes "
+                "in 1024x1024 resolution."
             ),
             "params": 81974227,
             "path": "segformer_b5",

keras_hub/src/models/seq_2_seq_lm_preprocessor.py

@@ -151,7 +151,7 @@ class Seq2SeqLMPreprocessor(Preprocessor):
         # `sequence_length` is an alias for `decoder_sequence_length`
         sequence_length=None,
     ):
-        """Convert encoder and decoder input strings to integer token inputs for generation.
+        """Convert input strings to integer token inputs for generation.
 
         Similar to calling the layer for training, this method takes in a dict
         containing `"encoder_text"` and `"decoder_text"`, with strings or tensor

keras_hub/src/models/stable_diffusion_3/mmdit.py

@@ -595,9 +595,28 @@ class MMDiTBlock(layers.Layer):
         self.context_block.build(context_shape, timestep_embedding_shape)
 
     def _compute_attention(self, query, key, value):
+        batch_size = ops.shape(query)[0]
+
+        # Use the fast path when `ops.dot_product_attention` and flash attention
+        # are available.
+        if hasattr(ops, "dot_product_attention") and hasattr(
+            keras.config, "is_flash_attention_enabled"
+        ):
+            # `ops.dot_product_attention` is slower than the vanilla
+            # implementation in the tensorflow backend.
+            encoded = ops.dot_product_attention(
+                query,
+                key,
+                value,
+                scale=self._inverse_sqrt_key_dim,
+                flash_attention=keras.config.is_flash_attention_enabled(),
+            )
+            return ops.reshape(
+                encoded, (batch_size, -1, self.num_heads * self.head_dim)
+            )
+
         # Ref: jax.nn.dot_product_attention
         # https://github.com/jax-ml/jax/blob/db89c245ac66911c98f265a05956fdfa4bc79d83/jax/_src/nn/functions.py#L846
-        batch_size = ops.shape(query)[0]
         logits = ops.einsum("BTNH,BSNH->BNTS", query, key)
         logits = ops.multiply(logits, self._inverse_sqrt_key_dim)
         probs = self.softmax(logits)

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_backbone.py

@@ -4,7 +4,7 @@ from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.backbone import Backbone
-from keras_hub.src.models.stable_diffusion_3.flow_match_euler_discrete_scheduler import (
+from keras_hub.src.models.stable_diffusion_3.flow_match_euler_discrete_scheduler import (  # noqa: E501
     FlowMatchEulerDiscreteScheduler,
 )
 from keras_hub.src.models.stable_diffusion_3.mmdit import MMDiT

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_image_to_image.py

@@ -2,10 +2,10 @@ from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.image_to_image import ImageToImage
-from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (  # noqa: E501
     StableDiffusion3Backbone,
 )
-from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (  # noqa: E501
     StableDiffusion3TextToImagePreprocessor,
 )
 
@@ -26,13 +26,17 @@ class StableDiffusion3ImageToImage(ImageToImage):
 
     Use `generate()` to do image generation.
     ```python
+    prompt = (
+        "Astronaut in a jungle, cold color palette, muted colors, "
+        "detailed, 8k"
+    )
     image_to_image = keras_hub.models.StableDiffusion3ImageToImage.from_preset(
         "stable_diffusion_3_medium", image_shape=(512, 512, 3)
     )
     image_to_image.generate(
         {
             "images": np.ones((512, 512, 3), dtype="float32"),
-            "prompts": "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+            "prompts": prompt,
         }
     )
 
@@ -40,7 +44,10 @@ class StableDiffusion3ImageToImage(ImageToImage):
     image_to_image.generate(
         {
             "images": np.ones((2, 512, 512, 3), dtype="float32"),
-            "prompts": ["cute wallpaper art of a cat", "cute wallpaper art of a dog"],
+            "prompts": [
+                "cute wallpaper art of a cat",
+                "cute wallpaper art of a dog",
+            ],
         }
     )
 
@@ -48,7 +55,7 @@ class StableDiffusion3ImageToImage(ImageToImage):
     image_to_image.generate(
         {
             "images": np.ones((512, 512, 3), dtype="float32"),
-            "prompts": "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+            "prompts": prompt,
         }
         num_steps=50,
         guidance_scale=5.0,
@@ -59,7 +66,7 @@ class StableDiffusion3ImageToImage(ImageToImage):
     text_to_image.generate(
         {
             "images": np.ones((512, 512, 3), dtype="float32"),
-            "prompts": "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+            "prompts": prompt,
             "negative_prompts": "green color",
         }
     )

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_inpaint.py

@@ -2,10 +2,10 @@ from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.inpaint import Inpaint
-from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (  # noqa: E501
     StableDiffusion3Backbone,
 )
-from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (  # noqa: E501
     StableDiffusion3TextToImagePreprocessor,
 )
 

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image.py

@@ -1,10 +1,10 @@
 from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (  # noqa: E501
     StableDiffusion3Backbone,
 )
-from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (  # noqa: E501
     StableDiffusion3TextToImagePreprocessor,
 )
 from keras_hub.src.models.text_to_image import TextToImage
@@ -46,9 +46,13 @@ class StableDiffusion3TextToImage(TextToImage):
     )
 
     # Generate with `negative_prompts`.
+    prompt = (
+        "Astronaut in a jungle, cold color palette, muted colors, "
+        "detailed, 8k"
+    )
     text_to_image.generate(
         {
-            "prompts": "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+            "prompts": prompt,
             "negative_prompts": "green color",
         }
     )

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image_preprocessor.py

@@ -3,7 +3,7 @@ from keras import layers
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.preprocessor import Preprocessor
-from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (  # noqa: E501
     StableDiffusion3Backbone,
 )
 

keras_hub/src/models/task.py

@@ -149,7 +149,8 @@ class Task(PipelineModel):
 
         This constructor can be called in one of two ways. Either from a task
         specific base class like `keras_hub.models.CausalLM.from_preset()`, or
-        from a model class like `keras_hub.models.BertTextClassifier.from_preset()`.
+        from a model class like
+        `keras_hub.models.BertTextClassifier.from_preset()`.
         If calling from the a base class, the subclass of the returning object
         will be inferred from the config in the preset directory.
 
@@ -294,7 +295,8 @@ class Task(PipelineModel):
             return "(" + ", ".join(highlighted) + ")"
 
         if self.preprocessor:
-            # Create a rich console for printing. Capture for non-interactive logging.
+            # Create a rich console for printing. Capture for non-interactive
+            # logging.
             if print_fn:
                 console = rich_console.Console(
                     highlight=False, force_terminal=False, color_system=None

keras_hub/src/models/text_classifier.py

@@ -21,8 +21,8 @@ class TextClassifier(Task):
     To fine-tune with `fit()`, pass a dataset containing tuples of `(x, y)`
     labels where `x` is a string and `y` is a integer from `[0, num_classes)`.
 
-    All `TextClassifier` tasks include a `from_preset()` constructor which can be
-    used to load a pre-trained config and weights.
+    All `TextClassifier` tasks include a `from_preset()` constructor which can
+    be used to load a pre-trained config and weights.
 
     Some, but not all, classification presets include classification head
     weights in a `task.weights.h5` file. For these presets, you can omit passing

keras_hub/src/models/text_to_image.py

@@ -262,9 +262,13 @@ class TextToImage(Task):
         pass `prompts` and `negative_prompts` as a dict:
 
         ```python
+        prompt = (
+            "Astronaut in a jungle, cold color palette, muted colors, "
+            "detailed, 8k"
+        )
         text_to_image.generate(
             {
-                "prompts": "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+                "prompts": prompt,
                 "negative_prompts": "green color",
             }
         )

keras_hub/src/models/vae/vae_layers.py

@@ -159,7 +159,6 @@ class ResNetBlock(keras.layers.Layer):
         data_format=None,
         **kwargs,
     ):
-
         super().__init__(**kwargs)
         data_format = standardize_data_format(data_format)
         channel_axis = -1 if data_format == "channels_last" else 1

keras_hub/src/models/vit/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.vit.vit_backbone import ViTBackbone
+from keras_hub.src.models.vit.vit_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, ViTBackbone)