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/mmdit_block.py

@@ -0,0 +1,317 @@
+# 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 math
+
+from keras import layers
+from keras import models
+from keras import ops
+
+from keras_hub.src.utils.keras_utils import gelu_approximate
+
+
+class DismantledBlock(layers.Layer):
+    def __init__(
+        self,
+        num_heads,
+        hidden_dim,
+        mlp_ratio=4.0,
+        use_projection=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.mlp_ratio = mlp_ratio
+        self.use_projection = use_projection
+
+        head_dim = hidden_dim // num_heads
+        self.head_dim = head_dim
+        mlp_hidden_dim = int(hidden_dim * mlp_ratio)
+        self.mlp_hidden_dim = mlp_hidden_dim
+        num_modulations = 6 if use_projection else 2
+        self.num_modulations = num_modulations
+
+        self.adaptive_norm_modulation = models.Sequential(
+            [
+                layers.Activation("silu", dtype=self.dtype_policy),
+                layers.Dense(
+                    num_modulations * hidden_dim, dtype=self.dtype_policy
+                ),
+            ],
+            name="adaptive_norm_modulation",
+        )
+        self.norm1 = layers.LayerNormalization(
+            epsilon=1e-6,
+            center=False,
+            scale=False,
+            dtype=self.dtype_policy,
+            name="norm1",
+        )
+        self.attention_qkv = layers.Dense(
+            hidden_dim * 3, dtype=self.dtype_policy, name="attention_qkv"
+        )
+        if use_projection:
+            self.attention_proj = layers.Dense(
+                hidden_dim, dtype=self.dtype_policy, name="attention_proj"
+            )
+            self.norm2 = layers.LayerNormalization(
+                epsilon=1e-6,
+                center=False,
+                scale=False,
+                dtype=self.dtype_policy,
+                name="norm2",
+            )
+            self.mlp = models.Sequential(
+                [
+                    layers.Dense(
+                        mlp_hidden_dim,
+                        activation=gelu_approximate,
+                        dtype=self.dtype_policy,
+                    ),
+                    layers.Dense(
+                        hidden_dim,
+                        dtype=self.dtype_policy,
+                    ),
+                ],
+                name="mlp",
+            )
+
+    def build(self, inputs_shape, timestep_embedding):
+        self.adaptive_norm_modulation.build(timestep_embedding)
+        self.attention_qkv.build(inputs_shape)
+        self.norm1.build(inputs_shape)
+        if self.use_projection:
+            self.attention_proj.build(inputs_shape)
+            self.norm2.build(inputs_shape)
+            self.mlp.build(inputs_shape)
+
+    def _modulate(self, inputs, shift, scale):
+        shift = ops.expand_dims(shift, axis=1)
+        scale = ops.expand_dims(scale, axis=1)
+        return ops.add(ops.multiply(inputs, ops.add(scale, 1.0)), shift)
+
+    def _compute_pre_attention(self, inputs, timestep_embedding, training=None):
+        batch_size = ops.shape(inputs)[0]
+        if self.use_projection:
+            modulation = self.adaptive_norm_modulation(
+                timestep_embedding, training=training
+            )
+            modulation = ops.reshape(
+                modulation, (batch_size, 6, self.hidden_dim)
+            )
+            (
+                shift_msa,
+                scale_msa,
+                gate_msa,
+                shift_mlp,
+                scale_mlp,
+                gate_mlp,
+            ) = ops.unstack(modulation, 6, axis=1)
+            qkv = self.attention_qkv(
+                self._modulate(self.norm1(inputs), shift_msa, scale_msa),
+                training=training,
+            )
+            qkv = ops.reshape(
+                qkv, (batch_size, -1, 3, self.num_heads, self.head_dim)
+            )
+            q, k, v = ops.unstack(qkv, 3, axis=2)
+            return (q, k, v), (inputs, gate_msa, shift_mlp, scale_mlp, gate_mlp)
+        else:
+            modulation = self.adaptive_norm_modulation(
+                timestep_embedding, training=training
+            )
+            modulation = ops.reshape(
+                modulation, (batch_size, 2, self.hidden_dim)
+            )
+            shift_msa, scale_msa = ops.unstack(modulation, 2, axis=1)
+            qkv = self.attention_qkv(
+                self._modulate(self.norm1(inputs), shift_msa, scale_msa),
+                training=training,
+            )
+            qkv = ops.reshape(
+                qkv, (batch_size, -1, 3, self.num_heads, self.head_dim)
+            )
+            q, k, v = ops.unstack(qkv, 3, axis=2)
+            return (q, k, v)
+
+    def _compute_post_attention(
+        self, inputs, inputs_intermediates, training=None
+    ):
+        x, gate_msa, shift_mlp, scale_mlp, gate_mlp = inputs_intermediates
+        attn = self.attention_proj(inputs, training=training)
+        x = ops.add(x, ops.multiply(ops.expand_dims(gate_msa, axis=1), attn))
+        x = ops.add(
+            x,
+            ops.multiply(
+                ops.expand_dims(gate_mlp, axis=1),
+                self.mlp(
+                    self._modulate(self.norm2(x), shift_mlp, scale_mlp),
+                    training=training,
+                ),
+            ),
+        )
+        return x
+
+    def call(
+        self,
+        inputs,
+        timestep_embedding=None,
+        inputs_intermediates=None,
+        pre_attention=True,
+        training=None,
+    ):
+        if pre_attention:
+            return self._compute_pre_attention(
+                inputs, timestep_embedding, training=training
+            )
+        else:
+            return self._compute_post_attention(
+                inputs, inputs_intermediates, training=training
+            )
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "mlp_ratio": self.mlp_ratio,
+                "use_projection": self.use_projection,
+            }
+        )
+        return config
+
+
+class MMDiTBlock(layers.Layer):
+    def __init__(
+        self,
+        num_heads,
+        hidden_dim,
+        mlp_ratio=4.0,
+        use_context_projection=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_heads = num_heads
+        self.hidden_dim = hidden_dim
+        self.mlp_ratio = mlp_ratio
+        self.use_context_projection = use_context_projection
+
+        head_dim = hidden_dim // num_heads
+        self.head_dim = head_dim
+        self._inverse_sqrt_key_dim = 1.0 / math.sqrt(head_dim)
+        self._dot_product_equation = "aecd,abcd->acbe"
+        self._combine_equation = "acbe,aecd->abcd"
+
+        self.x_block = DismantledBlock(
+            num_heads=num_heads,
+            hidden_dim=hidden_dim,
+            mlp_ratio=mlp_ratio,
+            use_projection=True,
+            dtype=self.dtype_policy,
+            name="x_block",
+        )
+        self.context_block = DismantledBlock(
+            num_heads=num_heads,
+            hidden_dim=hidden_dim,
+            mlp_ratio=mlp_ratio,
+            use_projection=use_context_projection,
+            dtype=self.dtype_policy,
+            name="context_block",
+        )
+
+    def build(self, inputs_shape, context_shape, timestep_embedding_shape):
+        self.x_block.build(inputs_shape, timestep_embedding_shape)
+        self.context_block.build(context_shape, timestep_embedding_shape)
+
+    def _compute_attention(self, query, key, value):
+        query = ops.multiply(
+            query, ops.cast(self._inverse_sqrt_key_dim, query.dtype)
+        )
+        attention_scores = ops.einsum(self._dot_product_equation, key, query)
+        attention_scores = ops.nn.softmax(attention_scores, axis=-1)
+        attention_output = ops.einsum(
+            self._combine_equation, attention_scores, value
+        )
+        batch_size = ops.shape(attention_output)[0]
+        attention_output = ops.reshape(
+            attention_output, (batch_size, -1, self.num_heads * self.head_dim)
+        )
+        return attention_output
+
+    def call(self, inputs, context, timestep_embedding, training=None):
+        # Compute pre-attention.
+        x = inputs
+        if self.use_context_projection:
+            context_qkv, context_intermediates = self.context_block(
+                context,
+                timestep_embedding=timestep_embedding,
+                training=training,
+            )
+        else:
+            context_qkv = self.context_block(
+                context,
+                timestep_embedding=timestep_embedding,
+                training=training,
+            )
+        context_len = ops.shape(context_qkv[0])[1]
+        x_qkv, x_intermediates = self.x_block(
+            x, timestep_embedding=timestep_embedding, training=training
+        )
+        q = ops.concatenate([context_qkv[0], x_qkv[0]], axis=1)
+        k = ops.concatenate([context_qkv[1], x_qkv[1]], axis=1)
+        v = ops.concatenate([context_qkv[2], x_qkv[2]], axis=1)
+
+        # Compute attention.
+        attention = self._compute_attention(q, k, v)
+        context_attention = attention[:, :context_len]
+        x_attention = attention[:, context_len:]
+
+        # Compute post-attention.
+        x = self.x_block(
+            x_attention,
+            inputs_intermediates=x_intermediates,
+            pre_attention=False,
+            training=training,
+        )
+        if self.use_context_projection:
+            context = self.context_block(
+                context_attention,
+                inputs_intermediates=context_intermediates,
+                pre_attention=False,
+                training=training,
+            )
+            return x, context
+        else:
+            return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_heads": self.num_heads,
+                "hidden_dim": self.hidden_dim,
+                "mlp_ratio": self.mlp_ratio,
+                "use_context_projection": self.use_context_projection,
+            }
+        )
+        return config
+
+    def compute_output_shape(
+        self, inputs_shape, context_shape, timestep_embedding_shape
+    ):
+        if self.use_context_projection:
+            return inputs_shape, context_shape
+        else:
+            return inputs_shape

keras_hub/src/models/stable_diffusion_v3/t5_xxl_preprocessor.py

@@ -0,0 +1,74 @@
+# 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.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.preprocessor import Preprocessor
+from keras_hub.src.models.t5.t5_tokenizer import T5Tokenizer
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+
+class T5XXLPreprocessor(Preprocessor):
+    tokenizer_cls = T5Tokenizer
+
+    def __init__(
+        self,
+        tokenizer,
+        sequence_length=256,
+        add_start_token=False,
+        add_end_token=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.tokenizer = tokenizer
+        self.sequence_length = sequence_length
+        self.add_start_token = add_start_token
+        self.add_end_token = add_end_token
+
+    def build(self, input_shape):
+        # Defer packer creation to `build()` so that we can be sure tokenizer
+        # assets have loaded when restoring a saved model.
+        self.packer = StartEndPacker(
+            start_value=self.tokenizer.start_token_id,
+            end_value=self.tokenizer.end_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            sequence_length=self.sequence_length,
+            return_padding_mask=True,
+        )
+        self.built = True
+
+    @preprocessing_function
+    def call(self, x, y=None, sample_weight=None, sequence_length=None):
+        token_ids, padding_mask = self.packer(
+            self.tokenizer(x),
+            sequence_length=sequence_length or self.sequence_length,
+            add_start_value=self.add_start_token,
+            add_end_value=self.add_end_token,
+        )
+        x = {
+            "token_ids": token_ids,
+            "padding_mask": padding_mask,
+        }
+        return keras.utils.pack_x_y_sample_weight(x, y, sample_weight)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "sequence_length": self.sequence_length,
+                "add_start_token": self.add_start_token,
+                "add_end_token": self.add_end_token,
+            }
+        )
+        return config

keras_hub/src/models/stable_diffusion_v3/t5_xxl_text_encoder.py

@@ -0,0 +1,155 @@
+# 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 (
+    ReversibleEmbedding,
+)
+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):
+    def __init__(
+        self,
+        vocabulary_size,
+        num_layers,
+        num_heads,
+        hidden_dim,
+        intermediate_dim,
+        key_value_dim=None,
+        dropout=0.1,
+        activation="relu",
+        use_gated_activation=True,
+        layer_norm_epsilon=1e-06,
+        tie_embedding_weights=True,
+        dtype=None,
+        **kwargs,
+    ):
+        # === Layers ===
+        self.token_embedding = ReversibleEmbedding(
+            input_dim=vocabulary_size,
+            output_dim=hidden_dim,
+            tie_weights=tie_embedding_weights,
+            embeddings_initializer=keras.initializers.TruncatedNormal(1.0),
+            dtype=dtype,
+            name="token_embedding",
+        )
+        self.encoder_embedding_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=dtype,
+            name="encoder_embedding_dropout",
+        )
+        self.encoder_transformer_layers = []
+        for i in range(num_layers):
+            layer = T5TransformerLayer(
+                is_decoder=False,
+                hidden_dim=hidden_dim,
+                intermediate_dim=intermediate_dim,
+                key_value_dim=key_value_dim or hidden_dim // num_heads,
+                dropout=dropout,
+                activation=activation,
+                layer_norm_epsilon=layer_norm_epsilon,
+                num_heads=num_heads,
+                use_gated_activation=use_gated_activation,
+                use_relative_attention_bias=bool(i == 0),
+                dtype=dtype,
+                name=f"transformer_encoder_layer_{i}",
+            )
+            self.encoder_transformer_layers.append(layer)
+        self.encoder_layer_norm = T5LayerNorm(
+            epsilon=layer_norm_epsilon,
+            dtype=dtype,
+            name="encoder_output_layer_norm",
+        )
+        self.encoder_dropout = keras.layers.Dropout(
+            dropout,
+            dtype=dtype,
+            name="encoder_output_dropout",
+        )
+
+        # === Functional Model ===
+        encoder_token_id_input = keras.Input(
+            shape=(None,), dtype="int32", name="encoder_token_ids"
+        )
+        encoder_padding_mask_input = keras.Input(
+            shape=(None,), dtype="int32", name="encoder_padding_mask"
+        )
+        # Encoder.
+        x = self.token_embedding(encoder_token_id_input)
+        x = self.encoder_embedding_dropout(x)
+        encoder_attention_mask = encoder_padding_mask_input[:, None, :]
+        position_bias = None
+        for transformer_layer in self.encoder_transformer_layers:
+            output = transformer_layer(
+                x,
+                attention_mask=encoder_attention_mask,
+                position_bias=position_bias,
+                use_causal_mask=False,
+            )
+            if isinstance(output, tuple):
+                x, position_bias = output
+        x = self.encoder_layer_norm(x)
+        x = self.encoder_dropout(x)
+        encoder_output = x
+
+        super().__init__(
+            {
+                "encoder_token_ids": encoder_token_id_input,
+                "encoder_padding_mask": encoder_padding_mask_input,
+            },
+            outputs=encoder_output,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.vocabulary_size = vocabulary_size
+        self.hidden_dim = hidden_dim
+        self.intermediate_dim = intermediate_dim
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.activation = keras.activations.get(activation)
+        self.key_value_dim = key_value_dim
+        self.dropout = dropout
+        self.use_gated_activation = use_gated_activation
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.tie_embedding_weights = tie_embedding_weights
+
+        if dtype is not None:
+            try:
+                self.dtype_policy = keras.dtype_policies.get(dtype)
+            # Before Keras 3.2, there is no `keras.dtype_policies.get`.
+            except AttributeError:
+                if isinstance(dtype, keras.DTypePolicy):
+                    dtype = dtype.name
+                self.dtype_policy = keras.DTypePolicy(dtype)
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "vocabulary_size": self.vocabulary_size,
+                "hidden_dim": self.hidden_dim,
+                "intermediate_dim": self.intermediate_dim,
+                "num_layers": self.num_layers,
+                "num_heads": self.num_heads,
+                "activation": keras.activations.serialize(self.activation),
+                "key_value_dim": self.key_value_dim,
+                "dropout": self.dropout,
+                "use_gated_activation": self.use_gated_activation,
+                "layer_norm_epsilon": self.layer_norm_epsilon,
+                "tie_embedding_weights": self.tie_embedding_weights,
+            }
+        )
+        return config

keras_hub/src/models/stable_diffusion_v3/vae_attention.py

@@ -0,0 +1,126 @@
+# 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 math
+
+from keras import layers
+from keras import ops
+
+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=self.dtype_policy,
+            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