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

keras_hub/src/models/sam/sam_transformer.py

@@ -0,0 +1,159 @@
+# 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
+
+from keras_hub.src.models.sam.sam_layers import (
+    MultiHeadAttentionWithDownsampling,
+)
+from keras_hub.src.models.sam.sam_layers import TwoWayMultiHeadAttention
+
+
+class TwoWayTransformer(keras.layers.Layer):
+    """A two-way cross-attention transformer decoder.
+
+    A transformer decoder that attends to an input image using
+    queries whose positional embedding is supplied.
+    The transformer decoder design is shown in
+    [1](https://arxiv.org/abs/2304.02643).
+    Each decoder layer performs 4 steps:
+    (1) self-attention on the tokens,
+    (2) cross-attention from tokens (as queries) to the image embedding,
+    (3) a point-wise MLPupdates each token, and
+    (4) cross-attention from the image embedding (as
+    queries) to tokens. This last step updates the image embedding with prompt
+    information. Each self/cross-attention and MLP has a residual connection
+    and layer normalization.
+    To ensure the decoder has access to critical geometric information the
+    positional encodings are added to the image embedding whenever they
+    participate in an attention layer. Additionally, the entire original
+    prompt tokens (including their positional encodings) are re-added to the
+    updated tokens whenever they participate in an attention layer. This
+    allows for a strong dependence on both the prompt token's geometric
+    location and type.
+
+    Args:
+        num_layers: int, optional. The num_layers of the attention blocks (the number
+            of attention blocks to use). Defaults to `2`.
+        hidden_size: int, optional. The number of features of the input image
+            and point embeddings. Defaults to `256`.
+        num_heads: int, optional. Number of heads to use in the attention
+            layers. Defaults to `8`.
+        intermediate_dim: int, optional. The number of units in the hidden layer of
+            the MLP block used in the attention layers. Defaults to `2048`.
+        activation: str, optional. The activation of the MLP block's output
+            layer used in the attention layers. Defaults to `"relu"`.
+        attention_downsample_rate: int, optional. The downsample rate of the
+            attention layers. Defaults to `2`.
+    """
+
+    def __init__(
+        self,
+        *,
+        num_layers=2,
+        hidden_size=256,
+        num_heads=8,
+        intermediate_dim=2048,
+        activation="relu",
+        attention_downsample_rate=2,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_layers = num_layers
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.intermediate_dim = intermediate_dim
+        self.activation = activation
+        self.attention_downsample_rate = attention_downsample_rate
+        self.layers = []
+        for i in range(num_layers):
+            self.layers.append(
+                TwoWayMultiHeadAttention(
+                    num_heads=num_heads,
+                    key_dim=hidden_size // num_heads,
+                    intermediate_dim=intermediate_dim,
+                    skip_first_layer_pos_embedding=(i == 0),
+                    attention_downsample_rate=attention_downsample_rate,
+                    activation=activation,
+                    dtype=self.dtype_policy,
+                )
+            )
+        self.final_attention_token_to_image = (
+            MultiHeadAttentionWithDownsampling(
+                num_heads=num_heads,
+                key_dim=hidden_size // num_heads,
+                downsample_rate=attention_downsample_rate,
+                dtype=self.dtype_policy,
+            )
+        )
+        self.final_layer_norm = keras.layers.LayerNormalization(
+            epsilon=1e-5, dtype=self.dtype_policy
+        )
+
+    def build(self, input_shape=None):
+        for layer in self.layers:
+            layer.build()
+        self.final_attention_token_to_image.build()
+        self.final_layer_norm.build([None, None, self.hidden_size])
+        self.built = True
+
+    def call(
+        self, image_embedding, image_positional_embeddings, point_embedding
+    ):
+        shape = ops.shape(image_embedding)
+        B, H, W, C = shape[0], shape[1], shape[2], shape[3]
+        image_embedding = ops.reshape(image_embedding, (B, H * W, C))
+
+        shape = ops.shape(image_positional_embeddings)
+        B, H, W, C = shape[0], shape[1], shape[2], shape[3]
+        image_positional_embeddings = ops.reshape(
+            image_positional_embeddings, (B, H * W, C)
+        )
+        queries = point_embedding
+        keys = image_embedding
+
+        for layer in self.layers:
+            queries, keys = layer(
+                queries=queries,
+                keys=keys,
+                query_pos_embedding=point_embedding,
+                key_pos_embedding=image_positional_embeddings,
+            )
+
+        queries_with_pos_embedding = queries + point_embedding
+        keys_with_pos_embedding = keys + image_positional_embeddings
+        attention_map = self.final_attention_token_to_image(
+            query=queries_with_pos_embedding,
+            key=keys_with_pos_embedding,
+            value=keys,
+        )
+        queries = queries + attention_map
+        queries = self.final_layer_norm(queries)
+
+        return queries, keys
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_layers": self.num_layers,
+                "hidden_size": self.hidden_size,
+                "num_heads": self.num_heads,
+                "intermediate_dim": self.intermediate_dim,
+                "activation": self.activation,
+                "attention_downsample_rate": self.attention_downsample_rate,
+            }
+        )
+        return config

keras_hub/src/models/stable_diffusion_3/__init__.py

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

keras_hub/src/models/stable_diffusion_3/flow_match_euler_discrete_scheduler.py

@@ -0,0 +1,93 @@
+# 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 import layers
+from keras import ops
+
+
+class FlowMatchEulerDiscreteScheduler(layers.Layer):
+    """Flow-matching sampling euler scheduler.
+
+    This layer is used to compute the discrete sigmas for the diffusion chain.
+    Typically, the sigma refers to the amount of noise added during the
+    diffusion process.
+
+    Args:
+        num_train_timesteps: int. The number of diffusion steps to train the
+            model.
+        shift: float. The shift value for the timestep schedule.
+        **kwargs: other keyword arguments passed to `keras.layers.Layer`,
+            including `name`, `dtype` etc.
+
+    Call arguments:
+        inputs: The current step of the diffusion process.
+        num_steps: The total number of steps in the diffusion process.
+
+    References:
+    - [Common Diffusion Noise Schedules and Sample Steps are Flawed](
+    https://arxiv.org/abs/2305.08891).
+    - [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](
+    https://arxiv.org/abs/2403.03206).
+    """
+
+    def __init__(self, num_train_timesteps=1000, shift=1.0, **kwargs):
+        super().__init__(**kwargs)
+        self.num_train_timesteps = int(num_train_timesteps)
+        self.shift = float(shift)
+
+        timesteps = ops.linspace(
+            1, num_train_timesteps, num_train_timesteps, dtype="float32"
+        )
+        timesteps = ops.flip(timesteps, axis=0)
+        sigmas = self._timestep_to_sigma(timesteps)
+
+        self.timesteps = ops.multiply(sigmas, num_train_timesteps)
+        self.sigma_min = sigmas[-1]
+        self.sigma_max = sigmas[0]
+
+    def _sigma_to_timestep(self, sigma):
+        return sigma * self.num_train_timesteps
+
+    def _timestep_to_sigma(self, timestep):
+        sigma = ops.divide(timestep, self.num_train_timesteps)
+        if self.shift != 1.0:
+            sigma = ops.divide(
+                ops.multiply(self.shift, sigma),
+                ops.add(1, ops.multiply(self.shift - 1.0, sigma)),
+            )
+        return sigma
+
+    def call(self, inputs, num_steps):
+        start = self._sigma_to_timestep(self.sigma_max)
+        end = self._sigma_to_timestep(self.sigma_min)
+        step_size = ops.divide(
+            ops.subtract(end, start), ops.subtract(num_steps, 1)
+        )
+        timestep = ops.add(start, ops.multiply(inputs, step_size))
+        sigma = ops.maximum(self._timestep_to_sigma(timestep), 0.0)
+        timestep = self._sigma_to_timestep(sigma)
+        return sigma, timestep
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "num_train_timesteps": self.num_train_timesteps,
+                "shift": self.shift,
+            }
+        )
+        return config
+
+    def compute_output_shape(self):
+        # Returns a tuple of (sigma, timestep).
+        return (None,), (None,)