keras-hub-nightly 0.19.0.dev202501080345__tar.gz → 0.19.0.dev202501090358__tar.gz

{keras_hub_nightly-0.19.0.dev202501080345 → keras_hub_nightly-0.19.0.dev202501090358}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: keras-hub-nightly
-Version: 0.19.0.dev202501080345
+Version: 0.19.0.dev202501090358
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team
@@ -31,6 +31,17 @@ Requires-Dist: tensorflow-text
 Provides-Extra: extras
 Requires-Dist: rouge-score; extra == "extras"
 Requires-Dist: sentencepiece; extra == "extras"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 # KerasHub: Multi-framework Pretrained Models
 [![](https://github.com/keras-team/keras-hub/workflows/Tests/badge.svg?branch=master)](https://github.com/keras-team/keras-hub/actions?query=workflow%3ATests+branch%3Amaster)

{keras_hub_nightly-0.19.0.dev202501080345 → keras_hub_nightly-0.19.0.dev202501090358}/keras_hub/src/models/stable_diffusion_3/mmdit.py

@@ -15,9 +15,8 @@ class AdaptiveLayerNormalization(layers.Layer):
 
     Args:
         embedding_dim: int. The size of each embedding vector.
-        residual_modulation: bool. Whether to output the modulation parameters
-            of the residual connection within the block of the diffusion
-            transformers. Defaults to `False`.
+        num_modulations: int. The number of the modulation parameters. The
+            available values are `2`, `6` and `9`. Defaults to `2`.
         **kwargs: other keyword arguments passed to `keras.layers.Layer`,
             including `name`, `dtype` etc.
 
@@ -28,11 +27,17 @@ class AdaptiveLayerNormalization(layers.Layer):
     https://arxiv.org/abs/2212.09748).
     """
 
-    def __init__(self, hidden_dim, residual_modulation=False, **kwargs):
+    def __init__(self, hidden_dim, num_modulations=2, **kwargs):
         super().__init__(**kwargs)
-        self.hidden_dim = int(hidden_dim)
-        self.residual_modulation = bool(residual_modulation)
-        num_modulations = 6 if self.residual_modulation else 2
+        hidden_dim = int(hidden_dim)
+        num_modulations = int(num_modulations)
+        if num_modulations not in (2, 6, 9):
+            raise ValueError(
+                "`num_modulations` must be `2`, `6` or `9`. "
+                f"Received: num_modulations={num_modulations}"
+            )
+        self.hidden_dim = hidden_dim
+        self.num_modulations = num_modulations
 
         self.silu = layers.Activation("silu", dtype=self.dtype_policy)
         self.dense = layers.Dense(
@@ -52,40 +57,84 @@ class AdaptiveLayerNormalization(layers.Layer):
         self.norm.build(inputs_shape)
 
     def call(self, inputs, embeddings, training=None):
-        x = inputs
+        hidden_states = inputs
         emb = self.dense(self.silu(embeddings), training=training)
-        if self.residual_modulation:
-            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
-                ops.split(emb, 6, axis=1)
-            )
+        if self.num_modulations == 9:
+            (
+                shift_msa,
+                scale_msa,
+                gate_msa,
+                shift_mlp,
+                scale_mlp,
+                gate_mlp,
+                shift_msa2,
+                scale_msa2,
+                gate_msa2,
+            ) = ops.split(emb, self.num_modulations, axis=1)
+        elif self.num_modulations == 6:
+            (
+                shift_msa,
+                scale_msa,
+                gate_msa,
+                shift_mlp,
+                scale_mlp,
+                gate_mlp,
+            ) = ops.split(emb, self.num_modulations, axis=1)
         else:
-            shift_msa, scale_msa = ops.split(emb, 2, axis=1)
+            shift_msa, scale_msa = ops.split(emb, self.num_modulations, axis=1)
+
         scale_msa = ops.expand_dims(scale_msa, axis=1)
         shift_msa = ops.expand_dims(shift_msa, axis=1)
-        x = ops.add(
-            ops.multiply(
-                self.norm(x, training=training),
-                ops.add(1.0, scale_msa),
-            ),
-            shift_msa,
+        norm_hidden_states = ops.cast(
+            self.norm(hidden_states, training=training), scale_msa.dtype
+        )
+        hidden_states = ops.add(
+            ops.multiply(norm_hidden_states, ops.add(1.0, scale_msa)), shift_msa
         )
-        if self.residual_modulation:
-            return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+        if self.num_modulations == 9:
+            scale_msa2 = ops.expand_dims(scale_msa2, axis=1)
+            shift_msa2 = ops.expand_dims(shift_msa2, axis=1)
+            hidden_states2 = ops.add(
+                ops.multiply(norm_hidden_states, ops.add(1.0, scale_msa2)),
+                shift_msa2,
+            )
+            return (
+                hidden_states,
+                gate_msa,
+                shift_mlp,
+                scale_mlp,
+                gate_mlp,
+                hidden_states2,
+                gate_msa2,
+            )
+        elif self.num_modulations == 6:
+            return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp
         else:
-            return x
+            return hidden_states
 
     def get_config(self):
         config = super().get_config()
         config.update(
             {
                 "hidden_dim": self.hidden_dim,
-                "residual_modulation": self.residual_modulation,
+                "num_modulations": self.num_modulations,
             }
         )
         return config
 
     def compute_output_shape(self, inputs_shape, embeddings_shape):
-        if self.residual_modulation:
+        if self.num_modulations == 9:
+            return (
+                inputs_shape,
+                embeddings_shape,
+                embeddings_shape,
+                embeddings_shape,
+                embeddings_shape,
+                inputs_shape,
+                embeddings_shape,
+            )
+        elif self.num_modulations == 6:
             return (
                 inputs_shape,
                 embeddings_shape,
@@ -345,6 +394,27 @@ class TimestepEmbedding(layers.Layer):
         return output_shape
 
 
+def get_qk_norm(qk_norm=None, q_norm_name="q_norm", k_norm_name="k_norm"):
+    """Helper function to instantiate `LayerNormalization` layers."""
+    q_norm = None
+    k_norm = None
+    if qk_norm is None:
+        pass
+    elif qk_norm == "rms_norm":
+        q_norm = layers.LayerNormalization(
+            epsilon=1e-6, rms_scaling=True, dtype="float32", name=q_norm_name
+        )
+        k_norm = layers.LayerNormalization(
+            epsilon=1e-6, rms_scaling=True, dtype="float32", name=k_norm_name
+        )
+    else:
+        raise NotImplementedError(
+            "Supported `qk_norm` are `'rms_norm'` and `None`. "
+            f"Received: qk_norm={qk_norm}."
+        )
+    return q_norm, k_norm
+
+
 class DismantledBlock(layers.Layer):
     """A dismantled block used to compute pre- and post-attention.
 
@@ -356,6 +426,8 @@ class DismantledBlock(layers.Layer):
             the end of the block.
         qk_norm: Optional str. Whether to normalize the query and key tensors.
             Available options are `None` and `"rms_norm"`. Defaults to `None`.
+        use_dual_attention: bool. Whether to use a dual attention in the
+            block. Defaults to `False`.
         **kwargs: other keyword arguments passed to `keras.layers.Layer`,
             including `name`, `dtype` etc.
     """
@@ -367,6 +439,7 @@ class DismantledBlock(layers.Layer):
         mlp_ratio=4.0,
         use_projection=True,
         qk_norm=None,
+        use_dual_attention=False,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -375,6 +448,7 @@ class DismantledBlock(layers.Layer):
         self.mlp_ratio = mlp_ratio
         self.use_projection = use_projection
         self.qk_norm = qk_norm
+        self.use_dual_attention = use_dual_attention
 
         head_dim = hidden_dim // num_heads
         self.head_dim = head_dim
@@ -384,7 +458,7 @@ class DismantledBlock(layers.Layer):
         if use_projection:
             self.ada_layer_norm = AdaptiveLayerNormalization(
                 hidden_dim,
-                residual_modulation=True,
+                num_modulations=9 if use_dual_attention else 6,
                 dtype=self.dtype_policy,
                 name="ada_layer_norm",
             )
@@ -395,18 +469,10 @@ class DismantledBlock(layers.Layer):
         self.attention_qkv = layers.Dense(
             hidden_dim * 3, dtype=self.dtype_policy, name="attention_qkv"
         )
-        if qk_norm is not None and qk_norm == "rms_norm":
-            self.q_norm = layers.LayerNormalization(
-                epsilon=1e-6, rms_scaling=True, dtype="float32", name="q_norm"
-            )
-            self.k_norm = layers.LayerNormalization(
-                epsilon=1e-6, rms_scaling=True, dtype="float32", name="q_norm"
-            )
-        elif qk_norm is not None:
-            raise NotImplementedError(
-                "Supported `qk_norm` are `'rms_norm'` and `None`. "
-                f"Received: qk_norm={qk_norm}."
-            )
+        q_norm, k_norm = get_qk_norm(qk_norm)
+        if q_norm is not None:
+            self.q_norm = q_norm
+            self.k_norm = k_norm
         if use_projection:
             self.attention_proj = layers.Dense(
                 hidden_dim, dtype=self.dtype_policy, name="attention_proj"
@@ -426,6 +492,19 @@ class DismantledBlock(layers.Layer):
                 name="mlp",
             )
 
+        if use_dual_attention:
+            self.attention_qkv2 = layers.Dense(
+                hidden_dim * 3, dtype=self.dtype_policy, name="attention_qkv2"
+            )
+            q_norm2, k_norm2 = get_qk_norm(qk_norm, "q_norm2", "k_norm2")
+            if q_norm is not None:
+                self.q_norm2 = q_norm2
+                self.k_norm2 = k_norm2
+            if use_projection:
+                self.attention_proj2 = layers.Dense(
+                    hidden_dim, dtype=self.dtype_policy, name="attention_proj2"
+                )
+
     def build(self, inputs_shape, timestep_embedding):
         self.ada_layer_norm.build(inputs_shape, timestep_embedding)
         self.attention_qkv.build(inputs_shape)
@@ -437,6 +516,13 @@ class DismantledBlock(layers.Layer):
             self.attention_proj.build(inputs_shape)
             self.norm2.build(inputs_shape)
             self.mlp.build(inputs_shape)
+        if self.use_dual_attention:
+            self.attention_qkv2.build(inputs_shape)
+            if self.qk_norm is not None:
+                self.q_norm2.build([None, None, self.num_heads, self.head_dim])
+                self.k_norm2.build([None, None, self.num_heads, self.head_dim])
+            if self.use_projection:
+                self.attention_proj2.build(inputs_shape)
 
     def _modulate(self, inputs, shift, scale):
         inputs = ops.cast(inputs, self.compute_dtype)
@@ -456,8 +542,12 @@ class DismantledBlock(layers.Layer):
             )
             q, k, v = ops.unstack(qkv, 3, axis=2)
             if self.qk_norm is not None:
-                q = self.q_norm(q, training=training)
-                k = self.k_norm(k, training=training)
+                q = ops.cast(
+                    self.q_norm(q, training=training), self.compute_dtype
+                )
+                k = ops.cast(
+                    self.k_norm(k, training=training), self.compute_dtype
+                )
             return (q, k, v), (inputs, gate_msa, shift_mlp, scale_mlp, gate_mlp)
         else:
             x = self.ada_layer_norm(
@@ -469,8 +559,12 @@ class DismantledBlock(layers.Layer):
             )
             q, k, v = ops.unstack(qkv, 3, axis=2)
             if self.qk_norm is not None:
-                q = self.q_norm(q, training=training)
-                k = self.k_norm(k, training=training)
+                q = ops.cast(
+                    self.q_norm(q, training=training), self.compute_dtype
+                )
+                k = ops.cast(
+                    self.k_norm(k, training=training), self.compute_dtype
+                )
             return (q, k, v)
 
     def _compute_post_attention(
@@ -495,22 +589,95 @@ class DismantledBlock(layers.Layer):
         )
         return x
 
+    def _compute_pre_attention_with_dual_attention(
+        self, inputs, timestep_embedding, training=None
+    ):
+        batch_size = ops.shape(inputs)[0]
+        x, gate_msa, shift_mlp, scale_mlp, gate_mlp, x2, gate_msa2 = (
+            self.ada_layer_norm(inputs, timestep_embedding, training=training)
+        )
+        # Compute the main attention
+        qkv = self.attention_qkv(x, 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)
+        if self.qk_norm is not None:
+            q = ops.cast(self.q_norm(q, training=training), self.compute_dtype)
+            k = ops.cast(self.k_norm(k, training=training), self.compute_dtype)
+        # Compute the dual attention
+        qkv2 = self.attention_qkv2(x2, training=training)
+        qkv2 = ops.reshape(
+            qkv2, (batch_size, -1, 3, self.num_heads, self.head_dim)
+        )
+        q2, k2, v2 = ops.unstack(qkv2, 3, axis=2)
+        if self.qk_norm is not None:
+            q2 = ops.cast(
+                self.q_norm2(q2, training=training), self.compute_dtype
+            )
+            k2 = ops.cast(
+                self.k_norm2(k2, training=training), self.compute_dtype
+            )
+        return (
+            (q, k, v),
+            (q2, k2, v2),
+            (inputs, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2),
+        )
+
+    def _compute_post_attention_with_dual_attention(
+        self, inputs, inputs2, inputs_intermediates, training=None
+    ):
+        x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2 = (
+            inputs_intermediates
+        )
+        gate_msa = ops.expand_dims(gate_msa, axis=1)
+        shift_mlp = ops.expand_dims(shift_mlp, axis=1)
+        scale_mlp = ops.expand_dims(scale_mlp, axis=1)
+        gate_mlp = ops.expand_dims(gate_mlp, axis=1)
+        gate_msa2 = ops.expand_dims(gate_msa2, axis=1)
+        attn = self.attention_proj(inputs, training=training)
+        x = ops.add(x, ops.multiply(gate_msa, attn))
+        attn2 = self.attention_proj2(inputs2, training=training)
+        x = ops.add(x, ops.multiply(gate_msa2, attn2))
+        x = ops.add(
+            x,
+            ops.multiply(
+                gate_mlp,
+                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,
+        inputs2=None,  # For the dual attention.
         pre_attention=True,
         training=None,
     ):
         if pre_attention:
-            return self._compute_pre_attention(
-                inputs, timestep_embedding, training=training
-            )
+            if self.use_dual_attention:
+                return self._compute_pre_attention_with_dual_attention(
+                    inputs, timestep_embedding, training=training
+                )
+            else:
+                return self._compute_pre_attention(
+                    inputs, timestep_embedding, training=training
+                )
         else:
-            return self._compute_post_attention(
-                inputs, inputs_intermediates, training=training
-            )
+            if self.use_dual_attention:
+                return self._compute_post_attention_with_dual_attention(
+                    inputs, inputs2, inputs_intermediates, training=training
+                )
+            else:
+                return self._compute_post_attention(
+                    inputs, inputs_intermediates, training=training
+                )
 
     def get_config(self):
         config = super().get_config()
@@ -521,6 +688,7 @@ class DismantledBlock(layers.Layer):
                 "mlp_ratio": self.mlp_ratio,
                 "use_projection": self.use_projection,
                 "qk_norm": self.qk_norm,
+                "use_dual_attention": self.use_dual_attention,
             }
         )
         return config
@@ -542,6 +710,8 @@ class MMDiTBlock(layers.Layer):
             layer at the end of the context block.
         qk_norm: Optional str. Whether to normalize the query and key tensors.
             Available options are `None` and `"rms_norm"`. Defaults to `None`.
+        use_dual_attention: bool. Whether to use a dual attention in the
+            block. Defaults to `False`.
         **kwargs: other keyword arguments passed to `keras.layers.Layer`,
             including `name`, `dtype` etc.
 
@@ -557,6 +727,7 @@ class MMDiTBlock(layers.Layer):
         mlp_ratio=4.0,
         use_context_projection=True,
         qk_norm=None,
+        use_dual_attention=False,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -565,6 +736,7 @@ class MMDiTBlock(layers.Layer):
         self.mlp_ratio = mlp_ratio
         self.use_context_projection = use_context_projection
         self.qk_norm = qk_norm
+        self.use_dual_attention = use_dual_attention
 
         head_dim = hidden_dim // num_heads
         self.head_dim = head_dim
@@ -576,6 +748,7 @@ class MMDiTBlock(layers.Layer):
             mlp_ratio=mlp_ratio,
             use_projection=True,
             qk_norm=qk_norm,
+            use_dual_attention=use_dual_attention,
             dtype=self.dtype_policy,
             name="x_block",
         )
@@ -602,8 +775,6 @@ class MMDiTBlock(layers.Layer):
         if hasattr(ops, "dot_product_attention") and hasattr(
             keras.config, "is_flash_attention_enabled"
         ):
-            # `ops.dot_product_attention` is slower than the vanilla
-            # implementation in the tensorflow backend.
             encoded = ops.dot_product_attention(
                 query,
                 key,
@@ -643,9 +814,14 @@ class MMDiTBlock(layers.Layer):
                 training=training,
             )
         context_len = ops.shape(context_qkv[0])[1]
-        x_qkv, x_intermediates = self.x_block(
-            x, timestep_embedding=timestep_embedding, training=training
-        )
+        if self.x_block.use_dual_attention:
+            x_qkv, x_qkv2, x_intermediates = self.x_block(
+                x, timestep_embedding=timestep_embedding, training=training
+            )
+        else:
+            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)
@@ -656,12 +832,23 @@ class MMDiTBlock(layers.Layer):
         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.x_block.use_dual_attention:
+            q2, k2, v2 = x_qkv2
+            x_attention2 = self._compute_attention(q2, k2, v2)
+            x = self.x_block(
+                x_attention,
+                inputs_intermediates=x_intermediates,
+                inputs2=x_attention2,
+                pre_attention=False,
+                training=training,
+            )
+        else:
+            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,
@@ -682,6 +869,7 @@ class MMDiTBlock(layers.Layer):
                 "mlp_ratio": self.mlp_ratio,
                 "use_context_projection": self.use_context_projection,
                 "qk_norm": self.qk_norm,
+                "use_dual_attention": self.use_dual_attention,
             }
         )
         return config
@@ -761,6 +949,9 @@ class MMDiT(Backbone):
         qk_norm: Optional str. Whether to normalize the query and key tensors in
             the intermediate blocks. Available options are `None` and
             `"rms_norm"`. Defaults to `None`.
+        dual_attention_indices: Optional tuple. Specifies the indices of
+            the blocks that serve as dual attention blocks. Typically, this is
+            for 3.5 version. Defaults to `None`.
         data_format: `None` or str. If specified, either `"channels_last"` or
             `"channels_first"`. The ordering of the dimensions in the
             inputs. `"channels_last"` corresponds to inputs with shape
@@ -786,6 +977,7 @@ class MMDiT(Backbone):
         context_shape=(None, 4096),
         pooled_projection_shape=(2048,),
         qk_norm=None,
+        dual_attention_indices=None,
         data_format=None,
         dtype=None,
         **kwargs,
@@ -799,6 +991,7 @@ class MMDiT(Backbone):
         image_width = latent_shape[1] // patch_size
         output_dim = latent_shape[-1]
         output_dim_in_final = patch_size**2 * output_dim
+        dual_attention_indices = dual_attention_indices or ()
         data_format = standardize_data_format(data_format)
         if data_format != "channels_last":
             raise NotImplementedError(
@@ -840,6 +1033,7 @@ class MMDiT(Backbone):
                 mlp_ratio,
                 use_context_projection=not (i == num_layers - 1),
                 qk_norm=qk_norm,
+                use_dual_attention=i in dual_attention_indices,
                 dtype=dtype,
                 name=f"joint_block_{i}",
             )
@@ -910,6 +1104,7 @@ class MMDiT(Backbone):
         self.context_shape = context_shape
         self.pooled_projection_shape = pooled_projection_shape
         self.qk_norm = qk_norm
+        self.dual_attention_indices = dual_attention_indices
 
     def get_config(self):
         config = super().get_config()
@@ -925,6 +1120,7 @@ class MMDiT(Backbone):
                 "context_shape": self.context_shape,
                 "pooled_projection_shape": self.pooled_projection_shape,
                 "qk_norm": self.qk_norm,
+                "dual_attention_indices": self.dual_attention_indices,
             }
         )
         return config

{keras_hub_nightly-0.19.0.dev202501080345 → keras_hub_nightly-0.19.0.dev202501090358}/keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_backbone.py

@@ -205,7 +205,10 @@ class StableDiffusion3Backbone(Backbone):
         mmdit_qk_norm: Optional str. Whether to normalize the query and key
             tensors for each transformer in MMDiT. Available options are `None`
             and `"rms_norm"`. Typically, this is set to `None` for 3.0 version
-            and to `"rms_norm" for 3.5 version.
+            and to `"rms_norm"` for 3.5 version.
+        mmdit_dual_attention_indices: Optional tuple. Specifies the indices of
+            the blocks that serve as dual attention blocks. Typically, this is
+            for 3.5 version. Defaults to `None`.
         vae: The VAE used for transformations between pixel space and latent
             space.
         clip_l: The CLIP text encoder for encoding the inputs.
@@ -253,6 +256,7 @@ class StableDiffusion3Backbone(Backbone):
         mmdit_depth=4,
         mmdit_position_size=192,
         mmdit_qk_norm=None,
+        mmdit_dual_attention_indices=None,
         vae=vae,
         clip_l=clip_l,
         clip_g=clip_g,
@@ -268,6 +272,7 @@ class StableDiffusion3Backbone(Backbone):
         mmdit_num_heads,
         mmdit_position_size,
         mmdit_qk_norm,
+        mmdit_dual_attention_indices,
         vae,
         clip_l,
         clip_g,
@@ -319,6 +324,7 @@ class StableDiffusion3Backbone(Backbone):
             context_shape=context_shape,
             pooled_projection_shape=pooled_projection_shape,
             qk_norm=mmdit_qk_norm,
+            dual_attention_indices=mmdit_dual_attention_indices,
             data_format=data_format,
             dtype=dtype,
             name="diffuser",
@@ -454,6 +460,7 @@ class StableDiffusion3Backbone(Backbone):
         self.mmdit_num_heads = mmdit_num_heads
         self.mmdit_position_size = mmdit_position_size
         self.mmdit_qk_norm = mmdit_qk_norm
+        self.mmdit_dual_attention_indices = mmdit_dual_attention_indices
         self.latent_channels = latent_channels
         self.output_channels = output_channels
         self.num_train_timesteps = num_train_timesteps
@@ -590,6 +597,9 @@ class StableDiffusion3Backbone(Backbone):
                 "mmdit_num_heads": self.mmdit_num_heads,
                 "mmdit_position_size": self.mmdit_position_size,
                 "mmdit_qk_norm": self.mmdit_qk_norm,
+                "mmdit_dual_attention_indices": (
+                    self.mmdit_dual_attention_indices
+                ),
                 "vae": layers.serialize(self.vae),
                 "clip_l": layers.serialize(self.clip_l),
                 "clip_g": layers.serialize(self.clip_g),
@@ -638,7 +648,10 @@ class StableDiffusion3Backbone(Backbone):
             )
 
         # To maintain backward compatibility, we need to ensure that
-        # `mmdit_qk_norm` is included in the config.
+        # `mmdit_qk_norm` and `mmdit_dual_attention_indices` is included in the
+        # config.
         if "mmdit_qk_norm" not in config:
             config["mmdit_qk_norm"] = None
+        if "mmdit_dual_attention_indices" not in config:
+            config["mmdit_dual_attention_indices"] = None
         return cls(**config)

{keras_hub_nightly-0.19.0.dev202501080345 → keras_hub_nightly-0.19.0.dev202501090358}/keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_presets.py

@@ -13,6 +13,18 @@ backbone_presets = {
         },
         "kaggle_handle": "kaggle://keras/stablediffusion3/keras/stable_diffusion_3_medium/4",
     },
+    "stable_diffusion_3.5_medium": {
+        "metadata": {
+            "description": (
+                "3 billion parameter, including CLIP L and CLIP G text "
+                "encoders, MMDiT-X generative model, and VAE autoencoder. "
+                "Developed by Stability AI."
+            ),
+            "params": 3371793763,
+            "path": "stable_diffusion_3",
+        },
+        "kaggle_handle": "kaggle://keras/stablediffusion3/keras/stable_diffusion_3.5_medium/1",
+    },
     "stable_diffusion_3.5_large": {
         "metadata": {
             "description": (

{keras_hub_nightly-0.19.0.dev202501080345 → keras_hub_nightly-0.19.0.dev202501090358}/keras_hub/src/version_utils.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 
 # Unique source of truth for the version number.
-__version__ = "0.19.0.dev202501080345"
+__version__ = "0.19.0.dev202501090358"
 
 
 @keras_hub_export("keras_hub.version")

{keras_hub_nightly-0.19.0.dev202501080345 → keras_hub_nightly-0.19.0.dev202501090358}/keras_hub_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: keras-hub-nightly
-Version: 0.19.0.dev202501080345
+Version: 0.19.0.dev202501090358
 Summary: Industry-strength Natural Language Processing extensions for Keras.
 Home-page: https://github.com/keras-team/keras-hub
 Author: Keras team
@@ -31,6 +31,17 @@ Requires-Dist: tensorflow-text
 Provides-Extra: extras
 Requires-Dist: rouge-score; extra == "extras"
 Requires-Dist: sentencepiece; extra == "extras"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 # KerasHub: Multi-framework Pretrained Models
 [![](https://github.com/keras-team/keras-hub/workflows/Tests/badge.svg?branch=master)](https://github.com/keras-team/keras-hub/actions?query=workflow%3ATests+branch%3Amaster)