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

keras_hub/src/models/stable_diffusion_v3/clip_preprocessor.py

@@ -1,93 +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_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
-from keras_hub.src.models.preprocessor import Preprocessor
-from keras_hub.src.models.stable_diffusion_v3.clip_tokenizer import (
-    CLIPTokenizer,
-)
-from keras_hub.src.utils.tensor_utils import preprocessing_function
-
-try:
-    import tensorflow as tf
-except ImportError:
-    tf = None
-
-
-class CLIPPreprocessor(Preprocessor):
-    tokenizer_cls = CLIPTokenizer
-
-    def __init__(
-        self,
-        tokenizer,
-        sequence_length=77,
-        add_start_token=True,
-        add_end_token=False,
-        to_lower=True,
-        pad_with_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
-        self.to_lower = to_lower
-        self.pad_with_end_token = pad_with_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.
-        pad_value = self.tokenizer.pad_token_id
-        if self.pad_with_end_token:
-            pad_value = self.tokenizer.end_token_id
-
-        self.packer = StartEndPacker(
-            start_value=self.tokenizer.start_token_id,
-            end_value=self.tokenizer.end_token_id,
-            pad_value=pad_value,
-            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):
-        if self.to_lower:
-            x = tf.strings.lower(x)
-        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,
-                "to_lower": self.to_lower,
-                "pad_with_end_token": self.pad_with_end_token,
-            }
-        )
-        return config

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