keras-hub-nightly 0.15.0.dev20240823171555__py3-none-any.whl → 0.16.0.dev2024092017__py3-none-any.whl

keras_hub/src/models/stable_diffusion_v3/vae_image_decoder.py

@@ -0,0 +1,186 @@
+# 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 import layers
+
+from keras_hub.src.models.stable_diffusion_v3.vae_attention import VAEAttention
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+class VAEImageDecoder(keras.Model):
+    def __init__(
+        self,
+        stackwise_num_filters,
+        stackwise_num_blocks,
+        output_channels=3,
+        latent_shape=(None, None, 16),
+        data_format=None,
+        dtype=None,
+        **kwargs,
+    ):
+        data_format = standardize_data_format(data_format)
+        gn_axis = -1 if data_format == "channels_last" else 1
+
+        # === Functional Model ===
+        latent_inputs = layers.Input(shape=latent_shape)
+
+        x = layers.Conv2D(
+            stackwise_num_filters[0],
+            3,
+            1,
+            padding="same",
+            data_format=data_format,
+            dtype=dtype,
+            name="input_projection",
+        )(latent_inputs)
+        x = apply_resnet_block(
+            x,
+            stackwise_num_filters[0],
+            data_format=data_format,
+            dtype=dtype,
+            name="input_block0",
+        )
+        x = VAEAttention(
+            stackwise_num_filters[0],
+            data_format=data_format,
+            dtype=dtype,
+            name="input_attention",
+        )(x)
+        x = apply_resnet_block(
+            x,
+            stackwise_num_filters[0],
+            data_format=data_format,
+            dtype=dtype,
+            name="input_block1",
+        )
+
+        # Stacks.
+        for i, filters in enumerate(stackwise_num_filters):
+            for j in range(stackwise_num_blocks[i]):
+                x = apply_resnet_block(
+                    x,
+                    filters,
+                    data_format=data_format,
+                    dtype=dtype,
+                    name=f"block{i}_{j}",
+                )
+            if i != len(stackwise_num_filters) - 1:
+                # No upsamling in the last blcok.
+                x = layers.UpSampling2D(
+                    2,
+                    data_format=data_format,
+                    dtype=dtype,
+                    name=f"upsample_{i}",
+                )(x)
+                x = layers.Conv2D(
+                    filters,
+                    3,
+                    1,
+                    padding="same",
+                    data_format=data_format,
+                    dtype=dtype,
+                    name=f"upsample_{i}_conv",
+                )(x)
+
+        # Ouput block.
+        x = layers.GroupNormalization(
+            groups=32,
+            axis=gn_axis,
+            epsilon=1e-6,
+            dtype=dtype,
+            name="output_norm",
+        )(x)
+        x = layers.Activation("swish", dtype=dtype, name="output_activation")(x)
+        image_outputs = layers.Conv2D(
+            output_channels,
+            3,
+            1,
+            padding="same",
+            data_format=data_format,
+            dtype=dtype,
+            name="output_projection",
+        )(x)
+        super().__init__(inputs=latent_inputs, outputs=image_outputs, **kwargs)
+
+        # === Config ===
+        self.stackwise_num_filters = stackwise_num_filters
+        self.stackwise_num_blocks = stackwise_num_blocks
+        self.output_channels = output_channels
+        self.latent_shape = latent_shape
+
+        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(
+            {
+                "stackwise_num_filters": self.stackwise_num_filters,
+                "stackwise_num_blocks": self.stackwise_num_blocks,
+                "output_channels": self.output_channels,
+                "image_shape": self.latent_shape,
+            }
+        )
+        return config
+
+
+def apply_resnet_block(x, filters, data_format=None, dtype=None, name=None):
+    data_format = standardize_data_format(data_format)
+    gn_axis = -1 if data_format == "channels_last" else 1
+    input_filters = x.shape[gn_axis]
+
+    residual = x
+    x = layers.GroupNormalization(
+        groups=32, axis=gn_axis, epsilon=1e-6, dtype=dtype, name=f"{name}_norm1"
+    )(x)
+    x = layers.Activation("swish", dtype=dtype)(x)
+    x = layers.Conv2D(
+        filters,
+        3,
+        1,
+        padding="same",
+        data_format=data_format,
+        dtype=dtype,
+        name=f"{name}_conv1",
+    )(x)
+    x = layers.GroupNormalization(
+        groups=32, axis=gn_axis, epsilon=1e-6, dtype=dtype, name=f"{name}_norm2"
+    )(x)
+    x = layers.Activation("swish")(x)
+    x = layers.Conv2D(
+        filters,
+        3,
+        1,
+        padding="same",
+        data_format=data_format,
+        dtype=dtype,
+        name=f"{name}_conv2",
+    )(x)
+    if input_filters != filters:
+        residual = layers.Conv2D(
+            filters,
+            1,
+            1,
+            data_format=data_format,
+            dtype=dtype,
+            name=f"{name}_residual_projection",
+        )(residual)
+    x = layers.Add(dtype=dtype)([residual, x])
+    return x

keras_hub/src/models/t5/__init__.py

@@ -14,7 +14,6 @@
 
 from keras_hub.src.models.t5.t5_backbone import T5Backbone
 from keras_hub.src.models.t5.t5_presets import backbone_presets
-from keras_hub.src.models.t5.t5_tokenizer import T5Tokenizer
 from keras_hub.src.utils.preset_utils import register_presets
 
-register_presets(backbone_presets, (T5Backbone, T5Tokenizer))
+register_presets(backbone_presets, T5Backbone)

keras_hub/src/models/t5/t5_tokenizer.py

@@ -13,12 +13,18 @@
 # limitations under the License.
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.t5.t5_backbone import T5Backbone
 from keras_hub.src.tokenizers.sentence_piece_tokenizer import (
     SentencePieceTokenizer,
 )
 
 
-@keras_hub_export("keras_hub.models.T5Tokenizer")
+@keras_hub_export(
+    [
+        "keras_hub.tokenizers.T5Tokenizer",
+        "keras_hub.models.T5Tokenizer",
+    ]
+)
 class T5Tokenizer(SentencePieceTokenizer):
     """T5 tokenizer layer based on SentencePiece.
 
@@ -74,27 +80,11 @@ class T5Tokenizer(SentencePieceTokenizer):
     ```
     """
 
-    def __init__(self, proto, **kwargs):
-        self.end_token = "</s>"
-        self.pad_token = "<pad>"
+    backbone_cls = T5Backbone
 
+    def __init__(self, proto, **kwargs):
+        # T5 uses the same start token as end token, i.e., "<\s>".
+        self._add_special_token("</s>", "end_token")
+        self._add_special_token("</s>", "start_token")
+        self._add_special_token("<pad>", "pad_token")
         super().__init__(proto=proto, **kwargs)
-
-    def set_proto(self, proto):
-        super().set_proto(proto)
-        if proto is not None:
-            for token in [self.end_token, self.pad_token]:
-                if token not in self.get_vocabulary():
-                    raise ValueError(
-                        f"Cannot find token `'{token}'` in the provided "
-                        f"`vocabulary`. Please provide `'{token}'` in your "
-                        "`vocabulary` or use a pretrained `vocabulary` name."
-                    )
-            self.end_token_id = self.token_to_id(self.end_token)
-            self.pad_token_id = self.token_to_id(self.pad_token)
-            # T5 uses the same start token as end token, i.e., "<\s>".
-            self.start_token_id = self.end_token_id
-        else:
-            self.end_token_id = None
-            self.pad_token_id = None
-            self.start_token_id = None

keras_hub/src/models/task.py

@@ -22,18 +22,11 @@ from rich import table as rich_table
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.utils.keras_utils import print_msg
 from keras_hub.src.utils.pipeline_model import PipelineModel
-from keras_hub.src.utils.preset_utils import CONFIG_FILE
-from keras_hub.src.utils.preset_utils import MODEL_WEIGHTS_FILE
 from keras_hub.src.utils.preset_utils import TASK_CONFIG_FILE
 from keras_hub.src.utils.preset_utils import TASK_WEIGHTS_FILE
-from keras_hub.src.utils.preset_utils import check_config_class
-from keras_hub.src.utils.preset_utils import check_file_exists
-from keras_hub.src.utils.preset_utils import check_format
-from keras_hub.src.utils.preset_utils import get_file
-from keras_hub.src.utils.preset_utils import jax_memory_cleanup
-from keras_hub.src.utils.preset_utils import list_presets
-from keras_hub.src.utils.preset_utils import list_subclasses
-from keras_hub.src.utils.preset_utils import load_serialized_object
+from keras_hub.src.utils.preset_utils import builtin_presets
+from keras_hub.src.utils.preset_utils import find_subclass
+from keras_hub.src.utils.preset_utils import get_preset_loader
 from keras_hub.src.utils.preset_utils import save_serialized_object
 from keras_hub.src.utils.python_utils import classproperty
 
@@ -56,12 +49,17 @@ class Task(PipelineModel):
     to load a pre-trained config and weights. Calling `from_preset()` on a task
     will automatically instantiate a `keras_hub.models.Backbone` and
     `keras_hub.models.Preprocessor`.
+
+    Args:
+        compile: boolean, defaults to `True`. If `True` will compile the model
+            with default parameters on construction. Model can still be
+            recompiled with a new loss, optimizer and metrics before training.
     """
 
     backbone_cls = None
     preprocessor_cls = None
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, compile=True, **kwargs):
         super().__init__(*args, **kwargs)
         self._functional_layer_ids = set(
             id(layer) for layer in self._flatten_layers()
@@ -69,6 +67,9 @@ class Task(PipelineModel):
         self._initialized = True
         if self.backbone is not None:
             self.dtype_policy = self._backbone.dtype_policy
+        if compile:
+            # Default compilation.
+            self.compile()
 
     def preprocess_samples(self, x, y=None, sample_weight=None):
         if self.preprocessor is not None:
@@ -131,13 +132,7 @@ class Task(PipelineModel):
     @classproperty
     def presets(cls):
         """List built-in presets for a `Task` subclass."""
-        presets = list_presets(cls)
-        # We can also load backbone presets.
-        if cls.backbone_cls is not None:
-            presets.update(cls.backbone_cls.presets)
-        for subclass in list_subclasses(cls):
-            presets.update(subclass.presets)
-        return presets
+        return builtin_presets(cls)
 
     @classmethod
     def from_preset(
@@ -149,10 +144,10 @@ class Task(PipelineModel):
         """Instantiate a `keras_hub.models.Task` from a model preset.
 
         A preset is a directory of configs, weights and other file assets used
-        to save and load a pre-trained model. The `preset` can be passed as a
+        to save and load a pre-trained model. The `preset` can be passed as
         one of:
 
-        1. a built in preset identifier like `'bert_base_en'`
+        1. a built-in preset identifier like `'bert_base_en'`
         2. a Kaggle Models handle like `'kaggle://user/bert/keras/bert_base_en'`
         3. a Hugging Face handle like `'hf://user/bert_base_en'`
         4. a path to a local preset directory like `'./bert_base_en'`
@@ -162,16 +157,16 @@ 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.BertClassifier.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.
 
         Args:
-            preset: string. A built in preset identifier, a Kaggle Models
+            preset: string. A built-in preset identifier, a Kaggle Models
                 handle, a Hugging Face handle, or a path to a local directory.
-            load_weights: bool. If `True`, the weights will be loaded into the
-                model architecture. If `False`, the weights will be randomly
-                initialized.
+            load_weights: bool. If `True`, saved weights will be loaded into
+                the model architecture. If `False`, all weights will be
+                randomly initialized.
 
         Examples:
         ```python
@@ -181,100 +176,37 @@ class Task(PipelineModel):
         )
 
         # Load a Bert classification task.
-        model = keras_hub.models.Classifier.from_preset(
+        model = keras_hub.models.TextClassifier.from_preset(
             "bert_base_en",
             num_classes=2,
         )
         ```
         """
-        format = check_format(preset)
-
-        if format == "transformers":
-            if cls.backbone_cls is None:
-                raise ValueError("Backbone class is None")
-            if cls.preprocessor_cls is None:
-                raise ValueError("Preprocessor class is None")
-
-            backbone = cls.backbone_cls.from_preset(preset)
-            preprocessor = cls.preprocessor_cls.from_preset(preset)
-            return cls(backbone=backbone, preprocessor=preprocessor, **kwargs)
-
         if cls == Task:
             raise ValueError(
                 "Do not call `Task.from_preset()` directly. Instead call a "
                 "particular task class, e.g. "
-                "`keras_hub.models.Classifier.from_preset()` or "
-                "`keras_hub.models.BertClassifier.from_preset()`."
-            )
-        if "backbone" in kwargs:
-            raise ValueError(
-                "You cannot pass a `backbone` argument to the `from_preset` "
-                f"method. Instead, call the {cls.__name__} default "
-                "constructor with a `backbone` argument. "
-                f"Received: backbone={kwargs['backbone']}."
+                "`keras_hub.models.TextClassifier.from_preset()`."
             )
 
-        # Check if we should load a `task.json` directly.
-        load_task_config = False
-        if check_file_exists(preset, TASK_CONFIG_FILE):
-            task_preset_cls = check_config_class(preset, TASK_CONFIG_FILE)
-            if issubclass(task_preset_cls, cls):
-                load_task_config = True
-        if load_task_config:
-            # Task case.
-            task_preset_cls = check_config_class(preset, TASK_CONFIG_FILE)
-            task = load_serialized_object(preset, TASK_CONFIG_FILE)
-            if load_weights:
-                jax_memory_cleanup(task)
-                if check_file_exists(preset, TASK_WEIGHTS_FILE):
-                    task.load_task_weights(get_file(preset, TASK_WEIGHTS_FILE))
-                task.backbone.load_weights(get_file(preset, MODEL_WEIGHTS_FILE))
-            task.preprocessor.tokenizer.load_preset_assets(preset)
-            return task
-
-        # Backbone case.
-        # If `task.json` doesn't exist or the task preset class is different
-        # from the calling class, create the task based on `config.json`.
-        backbone_preset_cls = check_config_class(preset, CONFIG_FILE)
-        if backbone_preset_cls is not cls.backbone_cls:
-            subclasses = list_subclasses(cls)
-            subclasses = tuple(
-                filter(
-                    lambda x: x.backbone_cls == backbone_preset_cls,
-                    subclasses,
-                )
-            )
-            if len(subclasses) == 0:
-                raise ValueError(
-                    f"No registered subclass of `{cls.__name__}` can load "
-                    f"a `{backbone_preset_cls.__name__}`."
-                )
-            if len(subclasses) > 1:
-                names = ", ".join(f"`{x.__name__}`" for x in subclasses)
-                raise ValueError(
-                    f"Ambiguous call to `{cls.__name__}.from_preset()`. "
-                    f"Found multiple possible subclasses {names}. "
-                    "Please call `from_preset` on a subclass directly."
-                )
-            cls = subclasses[0]
-        # Forward dtype to the backbone.
-        backbone_kwargs = {}
-        if "dtype" in kwargs:
-            backbone_kwargs = {"dtype": kwargs.pop("dtype")}
-        backbone = backbone_preset_cls.from_preset(
-            preset, load_weights=load_weights, **backbone_kwargs
-        )
-        if "preprocessor" in kwargs:
-            preprocessor = kwargs.pop("preprocessor")
-        else:
-            preprocessor = cls.preprocessor_cls.from_preset(preset)
-        return cls(backbone=backbone, preprocessor=preprocessor, **kwargs)
+        loader = get_preset_loader(preset)
+        backbone_cls = loader.check_backbone_class()
+        # Detect the correct subclass if we need to.
+        if cls.backbone_cls != backbone_cls:
+            cls = find_subclass(preset, cls, backbone_cls)
+        # Specifically for classifiers, we never load task weights if
+        # num_classes is supplied. We handle this in the task base class because
+        # it is the same logic for classifiers regardless of modality (text,
+        # images, audio).
+        load_task_weights = "num_classes" not in kwargs
+        return loader.load_task(cls, load_weights, load_task_weights, **kwargs)
 
     def load_task_weights(self, filepath):
         """Load only the tasks specific weights not in the backbone."""
         if not str(filepath).endswith(".weights.h5"):
             raise ValueError(
-                "The filename must end in `.weights.h5`. Received: filepath={filepath}"
+                "The filename must end in `.weights.h5`. "
+                f"Received: filepath={filepath}"
             )
         backbone_layer_ids = set(id(w) for w in self.backbone._flatten_layers())
         keras.saving.load_weights(
@@ -361,7 +293,9 @@ class Task(PipelineModel):
             print_fn = print_msg
 
         def highlight_number(x):
-            return f"[color(45)]{x}[/]" if x is None else f"[color(34)]{x}[/]"
+            if x is None:
+                f"[color(45)]{x}[/]"
+            return f"[color(34)]{x:,}[/]"  # Format number with commas.
 
         def highlight_symbol(x):
             return f"[color(33)]{x}[/]"
@@ -369,6 +303,10 @@ class Task(PipelineModel):
         def bold_text(x):
             return f"[bold]{x}[/]"
 
+        def highlight_shape(shape):
+            highlighted = [highlight_number(x) for x in shape]
+            return "(" + ", ".join(highlighted) + ")"
+
         if self.preprocessor:
             # Create a rich console for printing. Capture for non-interactive logging.
             if print_fn:
@@ -380,27 +318,44 @@ class Task(PipelineModel):
                 console = rich_console.Console(highlight=False)
 
             column_1 = rich_table.Column(
-                "Tokenizer (type)",
+                "Layer (type)",
                 justify="left",
-                width=int(0.5 * line_length),
+                width=int(0.6 * line_length),
             )
             column_2 = rich_table.Column(
-                "Vocab #",
+                "Config",
                 justify="right",
-                width=int(0.5 * line_length),
+                width=int(0.4 * line_length),
             )
             table = rich_table.Table(
                 column_1, column_2, width=line_length, show_lines=True
             )
+
+            def add_layer(layer, info):
+                layer_name = markup.escape(layer.name)
+                layer_class = highlight_symbol(
+                    markup.escape(layer.__class__.__name__)
+                )
+                table.add_row(
+                    f"{layer_name} ({layer_class})",
+                    info,
+                )
+
             tokenizer = self.preprocessor.tokenizer
-            tokenizer_name = markup.escape(tokenizer.name)
-            tokenizer_class = highlight_symbol(
-                markup.escape(tokenizer.__class__.__name__)
-            )
-            table.add_row(
-                f"{tokenizer_name} ({tokenizer_class})",
-                highlight_number(f"{tokenizer.vocabulary_size():,}"),
-            )
+            if tokenizer:
+                info = "Vocab size: "
+                info += highlight_number(tokenizer.vocabulary_size())
+                add_layer(tokenizer, info)
+            image_converter = self.preprocessor.image_converter
+            if image_converter:
+                info = "Image size: "
+                info += highlight_shape(image_converter.image_size())
+                add_layer(image_converter, info)
+            audio_converter = self.preprocessor.audio_converter
+            if audio_converter:
+                info = "Audio shape: "
+                info += highlight_shape(audio_converter.audio_shape())
+                add_layer(audio_converter, info)
 
             # Print the to the console.
             preprocessor_name = markup.escape(self.preprocessor.name)

keras_hub/src/models/{classifier.py → text_classifier.py}

@@ -17,25 +17,36 @@ from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.task import Task
 
 
-@keras_hub_export("keras_hub.models.Classifier")
-class Classifier(Task):
+@keras_hub_export(
+    [
+        "keras_hub.models.TextClassifier",
+        "keras_hub.models.Classifier",
+    ]
+)
+class TextClassifier(Task):
     """Base class for all classification tasks.
 
-    `Classifier` tasks wrap a `keras_hub.models.Backbone` and
+    `TextClassifier` tasks wrap a `keras_hub.models.Backbone` and
     a `keras_hub.models.Preprocessor` to create a model that can be used for
-    sequence classification. `Classifier` tasks take an additional
+    sequence classification. `TextClassifier` tasks take an additional
     `num_classes` argument, controlling the number of predicted output classes.
 
     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 `Classifier` tasks include a `from_preset()` constructor which can be
+    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
+    `num_classes` to restore the saved classification head. For all presets, if
+    `num_classes` is passed as a kwarg to `from_preset()`, the classification
+    head will be randomly initialized.
+
     Example:
     ```python
     # Load a BERT classifier with pre-trained weights.
-    classifier = keras_hub.models.Classifier.from_preset(
+    classifier = keras_hub.models.TextClassifier.from_preset(
         "bert_base_en",
         num_classes=2,
     )
@@ -52,11 +63,6 @@ class Classifier(Task):
     ```
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        # Default compilation.
-        self.compile()
-
     def compile(
         self,
         optimizer="auto",
@@ -65,9 +71,9 @@ class Classifier(Task):
         metrics="auto",
         **kwargs,
     ):
-        """Configures the `Classifier` task for training.
+        """Configures the `TextClassifier` task for training.
 
-        The `Classifier` task extends the default compilation signature of
+        The `TextClassifier` task extends the default compilation signature of
         `keras.Model.compile` with defaults for `optimizer`, `loss`, and
         `metrics`. To override these defaults, pass any value
         to these arguments during compilation.