keras-hub-nightly 0.16.1.dev202410200345__py3-none-any.whl → 0.19.0.dev202412070351__py3-none-any.whl

keras_hub/src/models/flux/flux_maths.py

@@ -0,0 +1,218 @@
+import keras
+from keras import ops
+
+
+class TimestepEmbedding(keras.layers.Layer):
+    """
+    Creates sinusoidal timestep embeddings.
+
+    Call arguments:
+        t: KerasTensor of shape (N,), representing N indices, one per batch element.
+            These values may be fractional.
+        dim: int. The dimension of the output.
+        max_period: int, optional. Controls the minimum frequency of the embeddings. Defaults to 10000.
+        time_factor: float, optional. A scaling factor applied to `t`. Defaults to 1000.0.
+
+    Returns:
+        KerasTensor: A tensor of shape (N, D) representing the positional embeddings,
+                     where N is the number of batch elements and D is the specified dimension `dim`.
+    """
+
+    def call(self, t, dim, max_period=10000, time_factor=1000.0):
+        t = time_factor * t
+        half_dim = dim // 2
+        freqs = ops.exp(
+            ops.cast(-ops.log(max_period), dtype=t.dtype)
+            * ops.arange(half_dim, dtype=t.dtype)
+            / half_dim
+        )
+        args = t[:, None] * freqs[None]
+        embedding = ops.concatenate([ops.cos(args), ops.sin(args)], axis=-1)
+
+        if dim % 2 != 0:
+            embedding = ops.concatenate(
+                [embedding, ops.zeros_like(embedding[:, :1])], axis=-1
+            )
+
+        return embedding
+
+
+class RotaryPositionalEmbedding(keras.layers.Layer):
+    """
+    Applies Rotary Positional Embedding (RoPE) to the input tensor.
+
+    Call arguments:
+        pos: KerasTensor. The positional tensor with shape (..., n, d).
+        dim: int. The embedding dimension, should be even.
+        theta: int. The base frequency.
+
+    Returns:
+        KerasTensor: The tensor with applied RoPE transformation.
+    """
+
+    def call(self, pos, dim, theta):
+        scale = ops.arange(0, dim, 2, dtype="float32") / dim
+        omega = 1.0 / (theta**scale)
+        out = ops.einsum("...n,d->...nd", pos, omega)
+        out = ops.stack(
+            [ops.cos(out), -ops.sin(out), ops.sin(out), ops.cos(out)], axis=-1
+        )
+        out = ops.reshape(out, ops.shape(out)[:-1] + (2, 2))
+        return ops.cast(out, dtype="float32")
+
+
+class ApplyRoPE(keras.layers.Layer):
+    """
+    Applies the RoPE transformation to the query and key tensors.
+
+    Call arguments:
+        xq: KerasTensor. The query tensor of shape (..., L, D).
+        xk: KerasTensor. The key tensor of shape (..., L, D).
+        freqs_cis: KerasTensor. The frequency complex numbers tensor with shape (..., 2).
+
+    Returns:
+        tuple[KerasTensor, KerasTensor]: The transformed query and key tensors.
+    """
+
+    def call(self, xq, xk, freqs_cis):
+        xq_ = ops.reshape(xq, (*ops.shape(xq)[:-1], -1, 1, 2))
+        xk_ = ops.reshape(xk, (*ops.shape(xk)[:-1], -1, 1, 2))
+
+        xq_out = (
+            freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+        )
+        xk_out = (
+            freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+        )
+
+        return ops.reshape(xq_out, ops.shape(xq)), ops.reshape(
+            xk_out, ops.shape(xk)
+        )
+
+
+class FluxRoPEAttention(keras.layers.Layer):
+    """
+    Computes the attention mechanism with the RoPE transformation applied to the query and key tensors.
+
+    Args:
+        dropout_p: float, optional. Dropout probability. Defaults to 0.0.
+        is_causal: bool, optional. If True, applies causal masking. Defaults to False.
+
+    Call arguments:
+        q: KerasTensor. Query tensor of shape (..., L, D).
+        k: KerasTensor. Key tensor of shape (..., S, D).
+        v: KerasTensor. Value tensor of shape (..., S, D).
+        positional_encoding: KerasTensor. Positional encoding tensor.
+
+    Returns:
+        KerasTensor: The resulting tensor from the attention mechanism.
+    """
+
+    def __init__(self, dropout_p=0.0, is_causal=False):
+        super(FluxRoPEAttention, self).__init__()
+        self.dropout_p = dropout_p
+        self.is_causal = is_causal
+
+    def call(self, q, k, v, positional_encoding):
+        # Apply the RoPE transformation
+        q, k = ApplyRoPE()(q, k, positional_encoding)
+
+        # Scaled dot-product attention
+        x = scaled_dot_product_attention(
+            q, k, v, dropout_p=self.dropout_p, is_causal=self.is_causal
+        )
+        x = ops.transpose(x, (0, 2, 1, 3))
+        b, l, h, d = ops.shape(x)
+        return ops.reshape(x, (b, l, h * d))
+
+
+# TODO: This is probably already implemented in several places, but is needed to ensure numeric equivalence to the original
+# implementation. It uses torch.functional.scaled_dot_product_attention() - do we have an equivalent already in Keras?
+def scaled_dot_product_attention(
+    query,
+    key,
+    value,
+    attn_mask=None,
+    dropout_p=0.0,
+    is_causal=False,
+    scale=None,
+):
+    """
+    Computes the scaled dot-product attention.
+
+    Args:
+        query: KerasTensor. Query tensor of shape (..., L, D).
+        key: KerasTensor. Key tensor of shape (..., S, D).
+        value: KerasTensor. Value tensor of shape (..., S, D).
+        attn_mask: KerasTensor, optional. Attention mask tensor. Defaults to None.
+        dropout_p: float, optional. Dropout probability. Defaults to 0.0.
+        is_causal: bool, optional. If True, applies causal masking. Defaults to False.
+        scale: float, optional. Scale factor for attention. Defaults to None.
+
+    Returns:
+        KerasTensor: The output tensor from the attention mechanism.
+    """
+    L, S = ops.shape(query)[-2], ops.shape(key)[-2]
+    scale_factor = (
+        1 / ops.sqrt(ops.cast(ops.shape(query)[-1], dtype=query.dtype))
+        if scale is None
+        else scale
+    )
+    attn_bias = ops.zeros((L, S), dtype=query.dtype)
+
+    if is_causal:
+        assert attn_mask is None
+        temp_mask = ops.ones((L, S), dtype=ops.bool)
+        temp_mask = ops.tril(temp_mask, diagonal=0)
+        attn_bias = ops.where(temp_mask, attn_bias, float("-inf"))
+
+    if attn_mask is not None:
+        if ops.shape(attn_mask)[-1] == 1:  # If the mask is 3D
+            attn_bias += attn_mask
+        else:
+            attn_bias = ops.where(attn_mask, attn_bias, float("-inf"))
+
+    # Compute attention weights
+    attn_weight = (
+        ops.matmul(query, ops.transpose(key, axes=[0, 1, 3, 2])) * scale_factor
+    )
+    attn_weight += attn_bias
+    attn_weight = keras.activations.softmax(attn_weight, axis=-1)
+
+    if dropout_p > 0.0:
+        attn_weight = keras.layers.Dropout(dropout_p)(
+            attn_weight, training=True
+        )
+
+    return ops.matmul(attn_weight, value)
+
+
+def rearrange_symbolic_tensors(qkv, K, H):
+    """
+    Splits the qkv tensor into query (q), key (k), and value (v) components.
+
+    Mimics rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=num_heads),
+    for graph-mode TensorFlow support when doing functional subclassing
+    models.
+
+    Arguments:
+        qkv: np.ndarray. Input tensor of shape (B, L, K*H*D).
+        K: int. Number of components (q, k, v).
+        H: int. Number of attention heads.
+
+    Returns:
+        tuple: q, k, v tensors of shape (B, H, L, D).
+    """
+    # Get the shape of qkv and calculate L and D
+    B, L, dim = ops.shape(qkv)
+    D = dim // (K * H)
+
+    # Reshape and transpose the qkv tensor
+    qkv_reshaped = ops.reshape(qkv, (B, L, K, H, D))
+    qkv_transposed = ops.transpose(qkv_reshaped, (2, 0, 3, 1, 4))
+
+    # Split q, k, v along the first dimension (K)
+    qkv_splits = ops.split(qkv_transposed, K, axis=0)
+    q, k, v = [ops.squeeze(split, 0) for split in qkv_splits]
+
+    return q, k, v

keras_hub/src/models/flux/flux_model.py

@@ -0,0 +1,231 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.backbone import Backbone
+from keras_hub.src.models.flux.flux_layers import DoubleStreamBlock
+from keras_hub.src.models.flux.flux_layers import EmbedND
+from keras_hub.src.models.flux.flux_layers import LastLayer
+from keras_hub.src.models.flux.flux_layers import MLPEmbedder
+from keras_hub.src.models.flux.flux_layers import SingleStreamBlock
+from keras_hub.src.models.flux.flux_maths import TimestepEmbedding
+
+
+@keras_hub_export("keras_hub.models.FluxBackbone")
+class FluxBackbone(Backbone):
+    """
+    Transformer model for flow matching on sequences.
+
+    The model processes image and text data with associated positional and timestep
+    embeddings, and optionally applies guidance embedding. Double-stream blocks
+    handle separate image and text streams, while single-stream blocks combine
+    these streams. Ported from: https://github.com/black-forest-labs/flux
+
+    Args:
+        input_channels: int. The number of input channels.
+        hidden_size: int. The hidden size of the transformer, must be divisible by `num_heads`.
+        mlp_ratio: float. The ratio of the MLP dimension to the hidden size.
+        num_heads: int. The number of attention heads.
+        depth: int. The number of double-stream blocks.
+        depth_single_blocks: int. The number of single-stream blocks.
+        axes_dim: list[int]. A list of dimensions for the positional embedding axes.
+        theta: int. The base frequency for positional embeddings.
+        use_bias: bool. Whether to apply bias to the query, key, and value projections.
+        guidance_embed: bool. If True, applies guidance embedding in the model.
+
+    Call arguments:
+        image: KerasTensor. Image input tensor of shape (N, L, D) where N is the batch size,
+                L is the sequence length, and D is the feature dimension.
+        image_ids: KerasTensor. Image ID input tensor of shape (N, L, D) corresponding
+                to the image sequences.
+        text: KerasTensor. Text input tensor of shape (N, L, D).
+        text_ids: KerasTensor. Text ID input tensor of shape (N, L, D) corresponding
+            to the text sequences.
+        timesteps: KerasTensor. Timestep tensor used to compute positional embeddings.
+        y: KerasTensor. Additional vector input, such as target values.
+        guidance: KerasTensor, optional. Guidance input tensor used
+            in guidance-embedded models.
+    Raises:
+        ValueError: If `hidden_size` is not divisible by `num_heads`, or if `sum(axes_dim)` is not equal to the
+                    positional embedding dimension.
+    """
+
+    def __init__(
+        self,
+        input_channels,
+        hidden_size,
+        mlp_ratio,
+        num_heads,
+        depth,
+        depth_single_blocks,
+        axes_dim,
+        theta,
+        use_bias,
+        guidance_embed=False,
+        # These will be inferred from the CLIP/T5 encoders later
+        image_shape=(None, 768, 3072),
+        text_shape=(None, 768, 3072),
+        image_ids_shape=(None, 768, 3072),
+        text_ids_shape=(None, 768, 3072),
+        y_shape=(None, 128),
+        **kwargs,
+    ):
+
+        # === Layers ===
+        self.positional_embedder = EmbedND(theta=theta, axes_dim=axes_dim)
+        self.image_input_embedder = keras.layers.Dense(
+            hidden_size, use_bias=True
+        )
+        self.time_input_embedder = MLPEmbedder(hidden_dim=hidden_size)
+        self.vector_embedder = MLPEmbedder(hidden_dim=hidden_size)
+        self.guidance_input_embedder = (
+            MLPEmbedder(hidden_dim=hidden_size)
+            if guidance_embed
+            else keras.layers.Identity()
+        )
+        self.text_input_embedder = keras.layers.Dense(hidden_size)
+
+        self.double_blocks = [
+            DoubleStreamBlock(
+                hidden_size,
+                num_heads,
+                mlp_ratio=mlp_ratio,
+                use_bias=use_bias,
+            )
+            for _ in range(depth)
+        ]
+
+        self.single_blocks = [
+            SingleStreamBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio)
+            for _ in range(depth_single_blocks)
+        ]
+
+        self.final_layer = LastLayer(hidden_size, 1, input_channels)
+        self.timestep_embedding = TimestepEmbedding()
+        self.guidance_embed = guidance_embed
+
+        # === Functional Model ===
+        image_input = keras.Input(shape=image_shape, name="image")
+        image_ids = keras.Input(shape=image_ids_shape, name="image_ids")
+        text_input = keras.Input(shape=text_shape, name="text")
+        text_ids = keras.Input(shape=text_ids_shape, name="text_ids")
+        y = keras.Input(shape=y_shape, name="y")
+        timesteps_input = keras.Input(shape=(), name="timesteps")
+        guidance_input = keras.Input(shape=(), name="guidance")
+
+        # running on sequences image
+        image = self.image_input_embedder(image_input)
+        modulation_encoding = self.time_input_embedder(
+            self.timestep_embedding(timesteps_input, dim=256)
+        )
+        if self.guidance_embed:
+            if guidance_input is None:
+                raise ValueError(
+                    "Didn't get guidance strength for guidance distilled model."
+                )
+            modulation_encoding = (
+                modulation_encoding
+                + self.guidance_input_embedder(
+                    self.timestep_embedding(guidance_input, dim=256)
+                )
+            )
+
+        modulation_encoding = modulation_encoding + self.vector_embedder(y)
+        text = self.text_input_embedder(text_input)
+
+        ids = keras.ops.concatenate((text_ids, image_ids), axis=1)
+        positional_encoding = self.positional_embedder(ids)
+
+        for block in self.double_blocks:
+            image, text = block(
+                image=image,
+                text=text,
+                modulation_encoding=modulation_encoding,
+                positional_encoding=positional_encoding,
+            )
+
+        image = keras.ops.concatenate((text, image), axis=1)
+        for block in self.single_blocks:
+            image = block(
+                image,
+                modulation_encoding=modulation_encoding,
+                positional_encoding=positional_encoding,
+            )
+        image = image[:, text.shape[1] :, ...]
+
+        image = self.final_layer(
+            image, modulation_encoding
+        )  # (N, T, patch_size ** 2 * output_channels)
+
+        super().__init__(
+            inputs={
+                "image": image_input,
+                "image_ids": image_ids,
+                "text": text_input,
+                "text_ids": text_ids,
+                "y": y,
+                "timesteps": timesteps_input,
+                "guidance": guidance_input,
+            },
+            outputs=image,
+            **kwargs,
+        )
+
+        # === Config ===
+        self.input_channels = input_channels
+        self.output_channels = self.input_channels
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.image_shape = image_shape
+        self.text_shape = text_shape
+        self.image_ids_shape = image_ids_shape
+        self.text_ids_shape = text_ids_shape
+        self.y_shape = y_shape
+        self.mlp_ratio = mlp_ratio
+        self.depth = depth
+        self.depth_single_blocks = depth_single_blocks
+        self.axes_dim = axes_dim
+        self.theta = theta
+        self.use_bias = use_bias
+
+    def get_config(self):
+        config = super().get_config()
+        config.update(
+            {
+                "input_channels": self.input_channels,
+                "hidden_size": self.hidden_size,
+                "mlp_ratio": self.mlp_ratio,
+                "num_heads": self.num_heads,
+                "depth": self.depth,
+                "depth_single_blocks": self.depth_single_blocks,
+                "axes_dim": self.axes_dim,
+                "theta": self.theta,
+                "use_bias": self.use_bias,
+                "guidance_embed": self.guidance_embed,
+                "image_shape": self.image_shape,
+                "text_shape": self.text_shape,
+                "image_ids_shape": self.image_ids_shape,
+                "text_ids_shape": self.text_ids_shape,
+                "y_shape": self.y_shape,
+            }
+        )
+        return config
+
+    def encode_text_step(self, token_ids, negative_token_ids):
+        raise NotImplementedError("Not implemented yet")
+
+        def encode(token_ids):
+            raise NotImplementedError("Not implemented yet")
+
+    def encode_image_step(self, images):
+        raise NotImplementedError("Not implemented yet")
+
+    def add_noise_step(self, latents, noises, step, num_steps):
+        raise NotImplementedError("Not implemented yet")
+
+    def denoise_step(
+        self,
+    ):
+        raise NotImplementedError("Not implemented yet")
+
+    def decode_step(self, latents):
+        raise NotImplementedError("Not implemented yet")

keras_hub/src/models/flux/flux_presets.py

@@ -0,0 +1,14 @@
+"""FLUX model preset configurations."""
+
+presets = {
+    "schnell": {
+        "metadata": {
+            "description": (
+                "A 12 billion parameter rectified flow transformer capable of generating images from text descriptions."
+            ),
+            "params": 124439808,
+            "path": "flux",
+        },
+        "kaggle_handle": "TBA",
+    },
+}

keras_hub/src/models/flux/flux_text_to_image.py

@@ -0,0 +1,142 @@
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.flux.flux_model import FluxBackbone
+from keras_hub.src.models.flux.flux_text_to_image_preprocessor import (
+    FluxTextToImagePreprocessor,
+)
+from keras_hub.src.models.text_to_image import TextToImage
+
+
+@keras_hub_export("keras_hub.models.FluxTextToImage")
+class FluxTextToImage(TextToImage):
+    """An end-to-end Flux 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.FluxBackbone` instance.
+        preprocessor: A
+            `keras_hub.models.FluxTextToImagePreprocessor` instance.
+
+    Examples:
+
+    Use `generate()` to do image generation.
+    ```python
+    text_to_image = keras_hub.models.FluxTextToImage.from_preset(
+        "TBA", 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 `guidance_scale`.
+    text_to_image.generate(
+        "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+        num_steps=50,
+        guidance_scale=5.0,
+    )
+
+    # Generate with `negative_prompts`.
+    text_to_image.generate(
+        {
+            "prompts": "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+            "negative_prompts": "green color",
+        }
+    )
+    ```
+    """
+
+    backbone_cls = FluxBackbone
+    preprocessor_cls = FluxTextToImagePreprocessor
+
+    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 " "`FluxTextToImage`."
+        )
+
+    def generate_step(
+        self,
+        latents,
+        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 pair of (batch_size, num_tokens) tensor containing the
+                tokens based on the input prompts and negative 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.
+        """
+        token_ids, negative_token_ids = token_ids
+
+        # Encode prompts.
+        embeddings = self.backbone.encode_text_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,
+        num_steps=28,
+        guidance_scale=7.0,
+        seed=None,
+    ):
+        return super().generate(
+            inputs,
+            num_steps=num_steps,
+            guidance_scale=guidance_scale,
+            seed=seed,
+        )

keras_hub/src/models/flux/flux_text_to_image_preprocessor.py

@@ -0,0 +1,73 @@
+import keras
+from keras import layers
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.flux.flux_model import FluxBackbone
+from keras_hub.src.models.preprocessor import Preprocessor
+
+
+@keras_hub_export("keras_hub.models.FluxTextToImagePreprocessor")
+class FluxTextToImagePreprocessor(Preprocessor):
+    """Flux text-to-image model preprocessor.
+
+    This preprocessing layer is meant for use with
+    `keras_hub.models.FluxTextToImagePreprocessor`.
+
+    For use with generation, the layer exposes one methods
+    `generate_preprocess()`.
+
+    Args:
+        clip_l_preprocessor: A `keras_hub.models.CLIPPreprocessor` instance.
+        t5_preprocessor: A optional `keras_hub.models.T5Preprocessor` instance.
+    """
+
+    backbone_cls = FluxBackbone
+
+    def __init__(
+        self,
+        clip_l_preprocessor,
+        t5_preprocessor=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.clip_l_preprocessor = clip_l_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"]
+        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
+                ),
+                "t5_preprocessor": layers.serialize(self.t5_preprocessor),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        for layer_name in (
+            "clip_l_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)