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

keras_hub/src/models/stable_diffusion_v3/mmdit_block.py

@@ -1,317 +0,0 @@
-# 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/vae_attention.py

@@ -1,126 +0,0 @@
-# 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

keras_hub/src/models/stable_diffusion_v3/vae_image_decoder.py

@@ -1,186 +0,0 @@
-# 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