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

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image.py

@@ -0,0 +1,138 @@
+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 (
+    StableDiffusion3Backbone,
+)
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (
+    StableDiffusion3TextToImagePreprocessor,
+)
+from keras_hub.src.models.text_to_image import TextToImage
+
+
+@keras_hub_export("keras_hub.models.StableDiffusion3TextToImage")
+class StableDiffusion3TextToImage(TextToImage):
+    """An end-to-end Stable Diffusion 3 model for text-to-image generation.
+
+    This model has a `generate()` method, which generates image based on a
+    prompt.
+
+    Args:
+        backbone: A `keras_hub.models.StableDiffusion3Backbone` instance.
+        preprocessor: A
+            `keras_hub.models.StableDiffusion3TextToImagePreprocessor` instance.
+
+    Examples:
+
+    Use `generate()` to do image generation.
+    ```python
+    text_to_image = keras_hub.models.StableDiffusion3TextToImage.from_preset(
+        "stable_diffusion_3_medium", height=512, width=512
+    )
+    text_to_image.generate(
+        "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
+    )
+
+    # Generate with batched prompts.
+    text_to_image.generate(
+        ["cute wallpaper art of a cat", "cute wallpaper art of a dog"]
+    )
+
+    # Generate with different `num_steps` and `classifier_free_guidance_scale`.
+    text_to_image.generate(
+        "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+        num_steps=50,
+        classifier_free_guidance_scale=5.0,
+    )
+    ```
+    """
+
+    backbone_cls = StableDiffusion3Backbone
+    preprocessor_cls = StableDiffusion3TextToImagePreprocessor
+
+    def __init__(
+        self,
+        backbone,
+        preprocessor,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+        # === Functional Model ===
+        inputs = backbone.input
+        outputs = backbone.output
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )
+
+    def fit(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Currently, `fit` is not supported for "
+            "`StableDiffusion3TextToImage`."
+        )
+
+    def generate_step(
+        self,
+        latents,
+        token_ids,
+        negative_token_ids,
+        num_steps,
+        guidance_scale,
+    ):
+        """A compilable generation function for batched of inputs.
+
+        This function represents the inner, XLA-compilable, generation function
+        for batched inputs.
+
+        Args:
+            latents: A (batch_size, height, width, channels) tensor
+                containing the latents to start generation from. Typically, this
+                tensor is sampled from the Gaussian distribution.
+            token_ids: A (batch_size, num_tokens) tensor containing the
+                tokens based on the input prompts.
+            negative_token_ids: A (batch_size, num_tokens) tensor
+                 containing the negative tokens based on the input prompts.
+            num_steps: int. The number of diffusion steps to take.
+            guidance_scale: float. The classifier free guidance scale defined in
+                [Classifier-Free Diffusion Guidance](
+                https://arxiv.org/abs/2207.12598). Higher scale encourages to
+                generate images that are closely linked to prompts, usually at
+                the expense of lower image quality.
+        """
+        # Encode inputs.
+        embeddings = self.backbone.encode_step(token_ids, negative_token_ids)
+
+        # Denoise.
+        def body_fun(step, latents):
+            return self.backbone.denoise_step(
+                latents,
+                embeddings,
+                step,
+                num_steps,
+                guidance_scale,
+            )
+
+        latents = ops.fori_loop(0, num_steps, body_fun, latents)
+
+        # Decode.
+        return self.backbone.decode_step(latents)
+
+    def generate(
+        self,
+        inputs,
+        negative_inputs=None,
+        num_steps=28,
+        guidance_scale=7.0,
+        seed=None,
+    ):
+        return super().generate(
+            inputs,
+            negative_inputs=negative_inputs,
+            num_steps=num_steps,
+            guidance_scale=guidance_scale,
+            seed=seed,
+        )

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image_preprocessor.py

@@ -0,0 +1,83 @@
+import keras
+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 (
+    StableDiffusion3Backbone,
+)
+
+
+@keras_hub_export("keras_hub.models.StableDiffusion3TextToImagePreprocessor")
+class StableDiffusion3TextToImagePreprocessor(Preprocessor):
+    """Stable Diffusion 3 text-to-image model preprocessor.
+
+    This preprocessing layer is meant for use with
+    `keras_hub.models.StableDiffusion3TextToImage`.
+
+    For use with generation, the layer exposes one methods
+    `generate_preprocess()`.
+
+    Args:
+        clip_l_preprocessor: A `keras_hub.models.CLIPPreprocessor` instance.
+        clip_g_preprocessor: A `keras_hub.models.CLIPPreprocessor` instance.
+        t5_preprocessor: A optional `keras_hub.models.T5Preprocessor` instance.
+    """
+
+    backbone_cls = StableDiffusion3Backbone
+
+    def __init__(
+        self,
+        clip_l_preprocessor,
+        clip_g_preprocessor,
+        t5_preprocessor=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.clip_l_preprocessor = clip_l_preprocessor
+        self.clip_g_preprocessor = clip_g_preprocessor
+        self.t5_preprocessor = t5_preprocessor
+
+    @property
+    def sequence_length(self):
+        """The padded length of model input sequences."""
+        return self.clip_l_preprocessor.sequence_length
+
+    def build(self, input_shape):
+        self.built = True
+
+    def generate_preprocess(self, x):
+        token_ids = {}
+        token_ids["clip_l"] = self.clip_l_preprocessor(x)["token_ids"]
+        token_ids["clip_g"] = self.clip_g_preprocessor(x)["token_ids"]
+        if self.t5_preprocessor is not None:
+            token_ids["t5"] = self.t5_preprocessor(x)["token_ids"]
+        return token_ids
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "clip_l_preprocessor": layers.serialize(
+                    self.clip_l_preprocessor
+                ),
+                "clip_g_preprocessor": layers.serialize(
+                    self.clip_g_preprocessor
+                ),
+                "t5_preprocessor": layers.serialize(self.t5_preprocessor),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        for layer_name in (
+            "clip_l_preprocessor",
+            "clip_g_preprocessor",
+            "t5_preprocessor",
+        ):
+            if layer_name in config and isinstance(config[layer_name], dict):
+                config[layer_name] = keras.layers.deserialize(
+                    config[layer_name]
+                )
+        return cls(**config)

keras_hub/src/models/{stable_diffusion_v3/t5_xxl_text_encoder.py → stable_diffusion_3/t5_encoder.py}

@@ -1,16 +1,3 @@
-# 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.layers.modeling.reversible_embedding import (
@@ -20,7 +7,7 @@ from keras_hub.src.models.t5.t5_layer_norm import T5LayerNorm
 from keras_hub.src.models.t5.t5_transformer_layer import T5TransformerLayer
 
 
-class T5XXLTextEncoder(keras.Model):
+class T5Encoder(keras.Model):
     def __init__(
         self,
         vocabulary_size,
@@ -81,10 +68,10 @@ class T5XXLTextEncoder(keras.Model):
 
         # === Functional Model ===
         encoder_token_id_input = keras.Input(
-            shape=(None,), dtype="int32", name="encoder_token_ids"
+            shape=(None,), dtype="int32", name="token_ids"
         )
         encoder_padding_mask_input = keras.Input(
-            shape=(None,), dtype="int32", name="encoder_padding_mask"
+            shape=(None,), dtype="int32", name="padding_mask"
         )
         # Encoder.
         x = self.token_embedding(encoder_token_id_input)
@@ -102,14 +89,14 @@ class T5XXLTextEncoder(keras.Model):
                 x, position_bias = output
         x = self.encoder_layer_norm(x)
         x = self.encoder_dropout(x)
-        encoder_output = x
+        sequence_output = x
 
         super().__init__(
             {
-                "encoder_token_ids": encoder_token_id_input,
-                "encoder_padding_mask": encoder_padding_mask_input,
+                "token_ids": encoder_token_id_input,
+                "padding_mask": encoder_padding_mask_input,
             },
-            outputs=encoder_output,
+            outputs=sequence_output,
             **kwargs,
         )
 

keras_hub/src/models/stable_diffusion_3/vae_image_decoder.py

@@ -0,0 +1,320 @@
+import math
+
+from keras import layers
+from keras import ops
+
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.utils.keras_utils import standardize_data_format
+
+
+class VAEAttention(layers.Layer):
+    def __init__(self, filters, groups=32, data_format=None, **kwargs):
+        super().__init__(**kwargs)
+        self.filters = filters
+        self.data_format = standardize_data_format(data_format)
+        gn_axis = -1 if self.data_format == "channels_last" else 1
+
+        self.group_norm = layers.GroupNormalization(
+            groups=groups,
+            axis=gn_axis,
+            epsilon=1e-6,
+            dtype="float32",
+            name="group_norm",
+        )
+        self.query_conv2d = layers.Conv2D(
+            filters,
+            1,
+            1,
+            data_format=self.data_format,
+            dtype=self.dtype_policy,
+            name="query_conv2d",
+        )
+        self.key_conv2d = layers.Conv2D(
+            filters,
+            1,
+            1,
+            data_format=self.data_format,
+            dtype=self.dtype_policy,
+            name="key_conv2d",
+        )
+        self.value_conv2d = layers.Conv2D(
+            filters,
+            1,
+            1,
+            data_format=self.data_format,
+            dtype=self.dtype_policy,
+            name="value_conv2d",
+        )
+        self.softmax = layers.Softmax(dtype="float32")
+        self.output_conv2d = layers.Conv2D(
+            filters,
+            1,
+            1,
+            data_format=self.data_format,
+            dtype=self.dtype_policy,
+            name="output_conv2d",
+        )
+
+        self.groups = groups
+        self._inverse_sqrt_filters = 1.0 / math.sqrt(float(filters))
+
+    def build(self, input_shape):
+        self.group_norm.build(input_shape)
+        self.query_conv2d.build(input_shape)
+        self.key_conv2d.build(input_shape)
+        self.value_conv2d.build(input_shape)
+        self.output_conv2d.build(input_shape)
+
+    def call(self, inputs, training=None):
+        x = self.group_norm(inputs)
+        query = self.query_conv2d(x)
+        key = self.key_conv2d(x)
+        value = self.value_conv2d(x)
+
+        if self.data_format == "channels_first":
+            query = ops.transpose(query, (0, 2, 3, 1))
+            key = ops.transpose(key, (0, 2, 3, 1))
+            value = ops.transpose(value, (0, 2, 3, 1))
+        shape = ops.shape(inputs)
+        b = shape[0]
+        query = ops.reshape(query, (b, -1, self.filters))
+        key = ops.reshape(key, (b, -1, self.filters))
+        value = ops.reshape(value, (b, -1, self.filters))
+
+        # Compute attention.
+        query = ops.multiply(
+            query, ops.cast(self._inverse_sqrt_filters, query.dtype)
+        )
+        # [B, H0 * W0, C], [B, H1 * W1, C] -> [B, H0 * W0, H1 * W1]
+        attention_scores = ops.einsum("abc,adc->abd", query, key)
+        attention_scores = ops.cast(
+            self.softmax(attention_scores), self.compute_dtype
+        )
+        # [B, H2 * W2, C], [B, H0 * W0, H1 * W1] -> [B, H1 * W1 ,C]
+        attention_output = ops.einsum("abc,adb->adc", value, attention_scores)
+        x = ops.reshape(attention_output, shape)
+
+        x = self.output_conv2d(x)
+        if self.data_format == "channels_first":
+            x = ops.transpose(x, (0, 3, 1, 2))
+        x = ops.add(x, inputs)
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "filters": self.filters,
+                "groups": self.groups,
+            }
+        )
+        return config
+
+    def compute_output_shape(self, input_shape):
+        return input_shape
+
+
+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="float32",
+        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="float32",
+        name=f"{name}_norm2",
+    )(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}_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
+
+
+class VAEImageDecoder(Backbone):
+    """Decoder for the VAE model used in Stable Diffusion 3.
+
+    Args:
+        stackwise_num_filters: list of ints. The number of filters for each
+            stack.
+        stackwise_num_blocks: list of ints. The number of blocks for each stack.
+        output_channels: int. The number of channels in the output.
+        latent_shape: tuple. The shape of the latent image.
+        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.
+    """
+
+    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="float32",
+            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
+
+    @property
+    def scaling_factor(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 1.5305
+
+    @property
+    def shift_factor(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 0.0609
+
+    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

keras_hub/src/models/t5/__init__.py

@@ -1,17 +1,3 @@
-# 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.
-
 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.utils.preset_utils import register_presets

keras_hub/src/models/t5/t5_backbone.py

@@ -1,17 +1,3 @@
-# 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

keras_hub/src/models/t5/t5_layer_norm.py

@@ -1,17 +1,3 @@
-# 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 ops
 

keras_hub/src/models/t5/t5_multi_head_attention.py

@@ -1,17 +1,3 @@
-# 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
 import numpy as np
 from keras import ops