keras-hub-nightly 0.16.1.dev202410020340__py3-none-any.whl → 0.19.0.dev202501260345__py3-none-any.whl

keras_hub/src/models/task.py

@@ -1,19 +1,20 @@
-import os
-
 import keras
 from rich import console as rich_console
 from rich import markup
 from rich import table as rich_table
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.audio_converter import AudioConverter
+from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.preprocessor import Preprocessor
+from keras_hub.src.tokenizers.tokenizer import Tokenizer
 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 TASK_CONFIG_FILE
-from keras_hub.src.utils.preset_utils import TASK_WEIGHTS_FILE
 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.preset_utils import get_preset_saver
 from keras_hub.src.utils.python_utils import classproperty
 
 
@@ -58,10 +59,15 @@ class Task(PipelineModel):
             self.compile()
 
     def preprocess_samples(self, x, y=None, sample_weight=None):
-        if self.preprocessor is not None:
+        # If `preprocessor` is `None`, return inputs unaltered.
+        if self.preprocessor is None:
+            return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
+        # If `preprocessor` is `Preprocessor` subclass, pass labels as a kwarg.
+        if isinstance(self.preprocessor, Preprocessor):
             return self.preprocessor(x, y=y, sample_weight=sample_weight)
-        else:
-            return super().preprocess_samples(x, y, sample_weight)
+        # For other layers and callable, do not pass the label.
+        x = self.preprocessor(x)
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
 
     def __setattr__(self, name, value):
         # Work around setattr issues for Keras 2 and Keras 3 torch backend.
@@ -143,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.
 
@@ -178,7 +185,10 @@ class Task(PipelineModel):
         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:
+        if (
+            issubclass(backbone_cls, Backbone)
+            and 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
@@ -232,17 +242,8 @@ class Task(PipelineModel):
         Args:
             preset_dir: The path to the local model preset directory.
         """
-        if self.preprocessor is None:
-            raise ValueError(
-                "Cannot save `task` to preset: `Preprocessor` is not initialized."
-            )
-
-        save_serialized_object(self, preset_dir, config_file=TASK_CONFIG_FILE)
-        if self.has_task_weights():
-            self.save_task_weights(os.path.join(preset_dir, TASK_WEIGHTS_FILE))
-
-        self.preprocessor.save_to_preset(preset_dir)
-        self.backbone.save_to_preset(preset_dir)
+        saver = get_preset_saver(preset_dir)
+        saver.save_task(self)
 
     @property
     def layers(self):
@@ -280,7 +281,7 @@ class Task(PipelineModel):
 
         def highlight_number(x):
             if x is None:
-                f"[color(45)]{x}[/]"
+                return f"[color(45)]{x}[/]"
             return f"[color(34)]{x:,}[/]"  # Format number with commas.
 
         def highlight_symbol(x):
@@ -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
@@ -327,24 +329,30 @@ class Task(PipelineModel):
                     info,
                 )
 
-            tokenizer = self.preprocessor.tokenizer
-            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)
+            # Since the preprocessor might be nested with multiple `Tokenizer`,
+            # `ImageConverter`, `AudioConverter` and even other `Preprocessor`
+            # instances, we should recursively iterate through them.
+            preprocessor = self.preprocessor
+            if preprocessor and isinstance(preprocessor, keras.Layer):
+                for layer in preprocessor._flatten_layers(include_self=False):
+                    if isinstance(layer, Tokenizer):
+                        info = "Vocab size: "
+                        info += highlight_number(layer.vocabulary_size())
+                        add_layer(layer, info)
+                    elif isinstance(layer, ImageConverter):
+                        info = "Image size: "
+                        image_size = layer.image_size
+                        if image_size is None:
+                            image_size = (None, None)
+                        info += highlight_shape(image_size)
+                        add_layer(layer, info)
+                    elif isinstance(layer, AudioConverter):
+                        info = "Audio shape: "
+                        info += highlight_shape(layer.audio_shape())
+                        add_layer(layer, info)
 
             # Print the to the console.
-            preprocessor_name = markup.escape(self.preprocessor.name)
+            preprocessor_name = markup.escape(preprocessor.name)
             console.print(bold_text(f'Preprocessor: "{preprocessor_name}"'))
             console.print(table)
 

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

@@ -56,6 +56,11 @@ class TextToImage(Task):
         # Default compilation.
         self.compile()
 
+    @property
+    def support_negative_prompts(self):
+        """Whether the model supports `negative_prompts` key in `generate()`."""
+        return bool(True)
+
     @property
     def latent_shape(self):
         return tuple(self.backbone.latent_shape)
@@ -171,9 +176,26 @@ class TextToImage(Task):
         This function converts all inputs to tensors, adds a batch dimension if
         necessary, and returns a iterable "dataset like" object (either an
         actual `tf.data.Dataset` or a list with a single batch element).
+
+        The input format must be one of the following:
+        - A single string
+        - A list of strings
+        - A dict with "prompts" and/or "negative_prompts" keys
+        - A tf.data.Dataset with "prompts" and/or "negative_prompts" keys
+
+        The output will be a dict with "prompts" and/or "negative_prompts" keys.
         """
         if tf and isinstance(inputs, tf.data.Dataset):
-            return inputs.as_numpy_iterator(), False
+            _inputs = {
+                "prompts": inputs.map(
+                    lambda x: x["prompts"]
+                ).as_numpy_iterator()
+            }
+            if self.support_negative_prompts:
+                _inputs["negative_prompts"] = inputs.map(
+                    lambda x: x["negative_prompts"]
+                ).as_numpy_iterator()
+            return _inputs, False
 
         def normalize(x):
             if isinstance(x, str):
@@ -182,13 +204,24 @@ class TextToImage(Task):
                 return x[tf.newaxis], True
             return x, False
 
+        def get_dummy_prompts(x):
+            dummy_prompts = [""] * len(x)
+            if tf and isinstance(x, tf.Tensor):
+                return tf.convert_to_tensor(dummy_prompts)
+            else:
+                return dummy_prompts
+
         if isinstance(inputs, dict):
             for key in inputs:
                 inputs[key], input_is_scalar = normalize(inputs[key])
         else:
             inputs, input_is_scalar = normalize(inputs)
+            inputs = {"prompts": inputs}
 
-        return inputs, input_is_scalar
+        if self.support_negative_prompts and "negative_prompts" not in inputs:
+            inputs["negative_prompts"] = get_dummy_prompts(inputs["prompts"])
+
+        return [inputs], input_is_scalar
 
     def _normalize_generate_outputs(self, outputs, input_is_scalar):
         """Normalize user output from the generate function.
@@ -199,12 +232,11 @@ class TextToImage(Task):
         """
 
         def normalize(x):
-            outputs = ops.clip(ops.divide(ops.add(x, 1.0), 2.0), 0.0, 1.0)
+            outputs = ops.concatenate(x, axis=0)
+            outputs = ops.clip(ops.divide(ops.add(outputs, 1.0), 2.0), 0.0, 1.0)
             outputs = ops.cast(ops.round(ops.multiply(outputs, 255.0)), "uint8")
-            outputs = ops.convert_to_numpy(outputs)
-            if input_is_scalar:
-                outputs = outputs[0]
-            return outputs
+            outputs = ops.squeeze(outputs, 0) if input_is_scalar else outputs
+            return ops.convert_to_numpy(outputs)
 
         if isinstance(outputs[0], dict):
             normalized = {}
@@ -216,33 +248,62 @@ class TextToImage(Task):
     def generate(
         self,
         inputs,
-        negative_inputs,
         num_steps,
-        guidance_scale,
+        guidance_scale=None,
         seed=None,
     ):
-        """Generate image based on the provided `inputs` and `negative_inputs`.
+        """Generate image based on the provided `inputs`.
+
+        Typically, `inputs` contains a text description (known as a prompt) used
+        to guide the image generation.
+
+        Some models support a `negative_prompts` key, which helps steer the
+        model away from generating certain styles and elements. To enable this,
+        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": prompt,
+                "negative_prompts": "green color",
+            }
+        )
+        ```
 
         If `inputs` are a `tf.data.Dataset`, outputs will be generated
         "batch-by-batch" and concatenated. Otherwise, all inputs will be
         processed as batches.
 
         Args:
-            inputs: python data, tensor data, or a `tf.data.Dataset`.
-            negative_inputs: python data, tensor data, or a `tf.data.Dataset`.
-                Unlike `inputs`, these are used as negative inputs to guide the
-                generation. If not provided, it defaults to `""` for each input
-                in `inputs`.
+            inputs: python data, tensor data, or a `tf.data.Dataset`. The format
+                must be one of the following:
+                - A single string
+                - A list of strings
+                - A dict with "prompts" and/or "negative_prompts" keys
+                - A `tf.data.Dataset` with "prompts" and/or "negative_prompts"
+                    keys
             num_steps: int. The number of diffusion steps to take.
-            guidance_scale: float. The classifier free guidance scale defined in
-                [Classifier-Free Diffusion Guidance](
+            guidance_scale: Optional float. The classifier free guidance scale
+                defined in [Classifier-Free Diffusion Guidance](
                 https://arxiv.org/abs/2207.12598). A higher scale encourages
                 generating images more closely related to the prompts, typically
-                at the cost of lower image quality.
+                at the cost of lower image quality. Note that some models don't
+                utilize classifier-free guidance.
             seed: optional int. Used as a random seed.
         """
+        num_steps = int(num_steps)
+        guidance_scale = (
+            float(guidance_scale) if guidance_scale is not None else None
+        )
         num_steps = ops.convert_to_tensor(num_steps, "int32")
-        guidance_scale = ops.convert_to_tensor(guidance_scale)
+        if guidance_scale is not None and guidance_scale > 1.0:
+            guidance_scale = ops.convert_to_tensor(guidance_scale)
+        else:
+            guidance_scale = None
 
         # Setup our three main passes.
         # 1. Preprocessing strings to dense integer tensors.
@@ -251,32 +312,36 @@ class TextToImage(Task):
         generate_function = self.make_generate_function()
 
         def preprocess(x):
-            return self.preprocessor.generate_preprocess(x)
+            if self.preprocessor is not None:
+                return self.preprocessor.generate_preprocess(x)
+            else:
+                return x
+
+        def generate(x):
+            token_ids = x[0] if self.support_negative_prompts else x
+
+            # Initialize latents.
+            if isinstance(token_ids, dict):
+                arbitrary_key = list(token_ids.keys())[0]
+                batch_size = ops.shape(token_ids[arbitrary_key])[0]
+            else:
+                batch_size = ops.shape(token_ids)[0]
+            latent_shape = (batch_size,) + self.latent_shape[1:]
+            latents = random.normal(latent_shape, dtype="float32", seed=seed)
+
+            return generate_function(latents, x, num_steps, guidance_scale)
 
         # Normalize and preprocess inputs.
         inputs, input_is_scalar = self._normalize_generate_inputs(inputs)
-        if negative_inputs is None:
-            negative_inputs = [""] * len(inputs)
-        negative_inputs, _ = self._normalize_generate_inputs(negative_inputs)
-
-        if self.preprocessor is not None:
-            inputs = preprocess(inputs)
-            negative_inputs = preprocess(negative_inputs)
-        if isinstance(inputs, dict):
-            batch_size = len(inputs[list(inputs.keys())[0]])
+        if self.support_negative_prompts:
+            token_ids = [preprocess(x["prompts"]) for x in inputs]
+            negative_token_ids = [
+                preprocess(x["negative_prompts"]) for x in inputs
+            ]
+            inputs = [x for x in zip(token_ids, negative_token_ids)]
         else:
-            batch_size = len(inputs)
-
-        # Initialize random latents.
-        latent_shape = (batch_size,) + self.latent_shape[1:]
-        latents = random.normal(latent_shape, dtype="float32", seed=seed)
+            inputs = [preprocess(x["prompts"]) for x in inputs]
 
         # Text-to-image.
-        outputs = generate_function(
-            latents,
-            inputs,
-            negative_inputs,
-            num_steps,
-            guidance_scale,
-        )
+        outputs = [generate(x) for x in inputs]
         return self._normalize_generate_outputs(outputs, input_is_scalar)

keras_hub/src/models/vae/__init__.py

@@ -0,0 +1 @@
+from keras_hub.src.models.vae.vae_backbone import VAEBackbone

keras_hub/src/models/vae/vae_backbone.py

@@ -0,0 +1,184 @@
+import keras
+
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.vae.vae_layers import (
+    DiagonalGaussianDistributionSampler,
+)
+from keras_hub.src.models.vae.vae_layers import VAEDecoder
+from keras_hub.src.models.vae.vae_layers import VAEEncoder
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+class VAEBackbone(Backbone):
+    """Variational Autoencoder(VAE) backbone used in latent diffusion models.
+
+    When encoding, this model generates mean and log variance of the input
+    images. When decoding, it reconstructs images from the latent space.
+
+    Args:
+        encoder_num_filters: list of ints. The number of filters for each
+            block in encoder.
+        encoder_num_blocks: list of ints. The number of blocks for each block in
+            encoder.
+        decoder_num_filters: list of ints. The number of filters for each
+            block in decoder.
+        decoder_num_blocks: list of ints. The number of blocks for each block in
+            decoder.
+        sampler_method: str. The method of the sampler for the intermediate
+            output. Available methods are `"sample"` and `"mode"`. `"sample"`
+            draws from the distribution using both the mean and log variance.
+            `"mode"` draws from the distribution using the mean only. Defaults
+            to `sample`.
+        input_channels: int. The number of channels in the input.
+        sample_channels: int. The number of channels in the sample. Typically,
+            this indicates the intermediate output of VAE, which is mean and
+            log variance.
+        output_channels: int. The number of channels in the output.
+        scale: float. The scaling factor applied to the latent space to ensure
+            it has unit variance during training of the diffusion model.
+            Defaults to `1.5305`, which is the value used in Stable Diffusion 3.
+        shift: float. The shift factor applied to the latent space to ensure it
+            has zero mean during training of the diffusion model. Defaults to
+            `0.0609`, which is the value used in Stable Diffusion 3.
+        data_format: `None` or str. If specified, either `"channels_last"` or
+            `"channels_first"`. The ordering of the dimensions in the
+            inputs. `"channels_last"` corresponds to inputs with shape
+            `(batch_size, height, width, channels)`
+            while `"channels_first"` corresponds to inputs with shape
+            `(batch_size, channels, height, width)`. It defaults to the
+            `image_data_format` value found in your Keras config file at
+            `~/.keras/keras.json`. If you never set it, then it will be
+            `"channels_last"`.
+        dtype: `None` or str or `keras.mixed_precision.DTypePolicy`. The dtype
+            to use for the model's computations and weights.
+
+    Example:
+    ```Python
+    backbone = VAEBackbone(
+        encoder_num_filters=[32, 32, 32, 32],
+        encoder_num_blocks=[1, 1, 1, 1],
+        decoder_num_filters=[32, 32, 32, 32],
+        decoder_num_blocks=[1, 1, 1, 1],
+    )
+    input_data = ops.ones((2, self.height, self.width, 3))
+    output = backbone(input_data)
+    ```
+    """
+
+    def __init__(
+        self,
+        encoder_num_filters,
+        encoder_num_blocks,
+        decoder_num_filters,
+        decoder_num_blocks,
+        sampler_method="sample",
+        input_channels=3,
+        sample_channels=32,
+        output_channels=3,
+        scale=1.5305,
+        shift=0.0609,
+        data_format=None,
+        dtype=None,
+        **kwargs,
+    ):
+        data_format = standardize_data_format(data_format)
+        if data_format == "channels_last":
+            image_shape = (None, None, input_channels)
+            channel_axis = -1
+        else:
+            image_shape = (input_channels, None, None)
+            channel_axis = 1
+
+        # === Layers ===
+        self.encoder = VAEEncoder(
+            encoder_num_filters,
+            encoder_num_blocks,
+            output_channels=sample_channels,
+            data_format=data_format,
+            dtype=dtype,
+            name="encoder",
+        )
+        # Use `sample()` to define the functional model.
+        self.distribution_sampler = DiagonalGaussianDistributionSampler(
+            method=sampler_method,
+            axis=channel_axis,
+            dtype=dtype,
+            name="distribution_sampler",
+        )
+        self.decoder = VAEDecoder(
+            decoder_num_filters,
+            decoder_num_blocks,
+            output_channels=output_channels,
+            data_format=data_format,
+            dtype=dtype,
+            name="decoder",
+        )
+
+        # === Functional Model ===
+        image_input = keras.Input(shape=image_shape)
+        sample = self.encoder(image_input)
+        latent = self.distribution_sampler(sample)
+        image_output = self.decoder(latent)
+        super().__init__(
+            inputs=image_input,
+            outputs=image_output,
+            dtype=dtype,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.encoder_num_filters = encoder_num_filters
+        self.encoder_num_blocks = encoder_num_blocks
+        self.decoder_num_filters = decoder_num_filters
+        self.decoder_num_blocks = decoder_num_blocks
+        self.sampler_method = sampler_method
+        self.input_channels = input_channels
+        self.sample_channels = sample_channels
+        self.output_channels = output_channels
+        self._scale = scale
+        self._shift = shift
+
+    @property
+    def scale(self):
+        """The scaling factor for the latent space.
+
+        This is used to scale the latent space to have unit variance when
+        training the diffusion model.
+        """
+        return self._scale
+
+    @property
+    def shift(self):
+        """The shift factor for the latent space.
+
+        This is used to shift the latent space to have zero mean when
+        training the diffusion model.
+        """
+        return self._shift
+
+    def encode(self, inputs, **kwargs):
+        """Encode the input images into latent space."""
+        sample = self.encoder(inputs, **kwargs)
+        return self.distribution_sampler(sample)
+
+    def decode(self, inputs, **kwargs):
+        """Decode the input latent space into images."""
+        return self.decoder(inputs, **kwargs)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "encoder_num_filters": self.encoder_num_filters,
+                "encoder_num_blocks": self.encoder_num_blocks,
+                "decoder_num_filters": self.decoder_num_filters,
+                "decoder_num_blocks": self.decoder_num_blocks,
+                "sampler_method": self.sampler_method,
+                "input_channels": self.input_channels,
+                "sample_channels": self.sample_channels,
+                "output_channels": self.output_channels,
+                "scale": self.scale,
+                "shift": self.shift,
+            }
+        )
+        return config