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

keras_hub/src/models/stable_diffusion_3/mmdit.py

@@ -354,6 +354,8 @@ class DismantledBlock(layers.Layer):
         mlp_ratio: float. The expansion ratio of `MLP`.
         use_projection: bool. Whether to use an attention projection layer at
             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`.
         **kwargs: other keyword arguments passed to `keras.layers.Layer`,
             including `name`, `dtype` etc.
     """
@@ -364,6 +366,7 @@ class DismantledBlock(layers.Layer):
         hidden_dim,
         mlp_ratio=4.0,
         use_projection=True,
+        qk_norm=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -371,6 +374,7 @@ class DismantledBlock(layers.Layer):
         self.hidden_dim = hidden_dim
         self.mlp_ratio = mlp_ratio
         self.use_projection = use_projection
+        self.qk_norm = qk_norm
 
         head_dim = hidden_dim // num_heads
         self.head_dim = head_dim
@@ -391,6 +395,18 @@ 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}."
+            )
         if use_projection:
             self.attention_proj = layers.Dense(
                 hidden_dim, dtype=self.dtype_policy, name="attention_proj"
@@ -413,6 +429,10 @@ class DismantledBlock(layers.Layer):
     def build(self, inputs_shape, timestep_embedding):
         self.ada_layer_norm.build(inputs_shape, timestep_embedding)
         self.attention_qkv.build(inputs_shape)
+        if self.qk_norm is not None:
+            # [batch_size, sequence_length, num_heads, head_dim]
+            self.q_norm.build([None, None, self.num_heads, self.head_dim])
+            self.k_norm.build([None, None, self.num_heads, self.head_dim])
         if self.use_projection:
             self.attention_proj.build(inputs_shape)
             self.norm2.build(inputs_shape)
@@ -435,6 +455,9 @@ class DismantledBlock(layers.Layer):
                 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 = self.q_norm(q, training=training)
+                k = self.k_norm(k, training=training)
             return (q, k, v), (inputs, gate_msa, shift_mlp, scale_mlp, gate_mlp)
         else:
             x = self.ada_layer_norm(
@@ -445,6 +468,9 @@ class DismantledBlock(layers.Layer):
                 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 = self.q_norm(q, training=training)
+                k = self.k_norm(k, training=training)
             return (q, k, v)
 
     def _compute_post_attention(
@@ -494,6 +520,7 @@ class DismantledBlock(layers.Layer):
                 "hidden_dim": self.hidden_dim,
                 "mlp_ratio": self.mlp_ratio,
                 "use_projection": self.use_projection,
+                "qk_norm": self.qk_norm,
             }
         )
         return config
@@ -513,6 +540,8 @@ class MMDiTBlock(layers.Layer):
         mlp_ratio: float. The expansion ratio of `MLP`.
         use_context_projection: bool. Whether to use an attention projection
             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`.
         **kwargs: other keyword arguments passed to `keras.layers.Layer`,
             including `name`, `dtype` etc.
 
@@ -527,6 +556,7 @@ class MMDiTBlock(layers.Layer):
         hidden_dim,
         mlp_ratio=4.0,
         use_context_projection=True,
+        qk_norm=None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -534,6 +564,7 @@ class MMDiTBlock(layers.Layer):
         self.hidden_dim = hidden_dim
         self.mlp_ratio = mlp_ratio
         self.use_context_projection = use_context_projection
+        self.qk_norm = qk_norm
 
         head_dim = hidden_dim // num_heads
         self.head_dim = head_dim
@@ -544,6 +575,7 @@ class MMDiTBlock(layers.Layer):
             hidden_dim=hidden_dim,
             mlp_ratio=mlp_ratio,
             use_projection=True,
+            qk_norm=qk_norm,
             dtype=self.dtype_policy,
             name="x_block",
         )
@@ -552,6 +584,7 @@ class MMDiTBlock(layers.Layer):
             hidden_dim=hidden_dim,
             mlp_ratio=mlp_ratio,
             use_projection=use_context_projection,
+            qk_norm=qk_norm,
             dtype=self.dtype_policy,
             name="context_block",
         )
@@ -629,6 +662,7 @@ class MMDiTBlock(layers.Layer):
                 "hidden_dim": self.hidden_dim,
                 "mlp_ratio": self.mlp_ratio,
                 "use_context_projection": self.use_context_projection,
+                "qk_norm": self.qk_norm,
             }
         )
         return config
@@ -705,6 +739,9 @@ class MMDiT(Backbone):
         latent_shape: tuple. The shape of the latent image.
         context_shape: tuple. The shape of the context.
         pooled_projection_shape: tuple. The shape of the pooled projection.
+        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`.
         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
@@ -729,6 +766,7 @@ class MMDiT(Backbone):
         latent_shape=(64, 64, 16),
         context_shape=(None, 4096),
         pooled_projection_shape=(2048,),
+        qk_norm=None,
         data_format=None,
         dtype=None,
         **kwargs,
@@ -782,6 +820,7 @@ class MMDiT(Backbone):
                 hidden_dim,
                 mlp_ratio,
                 use_context_projection=not (i == num_layers - 1),
+                qk_norm=qk_norm,
                 dtype=dtype,
                 name=f"joint_block_{i}",
             )
@@ -851,6 +890,7 @@ class MMDiT(Backbone):
         self.latent_shape = latent_shape
         self.context_shape = context_shape
         self.pooled_projection_shape = pooled_projection_shape
+        self.qk_norm = qk_norm
 
     def get_config(self):
         config = super().get_config()
@@ -865,6 +905,7 @@ class MMDiT(Backbone):
                 "latent_shape": self.latent_shape,
                 "context_shape": self.context_shape,
                 "pooled_projection_shape": self.pooled_projection_shape,
+                "qk_norm": self.qk_norm,
             }
         )
         return config

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_backbone.py

@@ -202,6 +202,10 @@ class StableDiffusion3Backbone(Backbone):
             transformer in MMDiT.
         mmdit_position_size: int. The size of the height and width for the
             position embedding in MMDiT.
+        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.
         vae: The VAE used for transformations between pixel space and latent
             space.
         clip_l: The CLIP text encoder for encoding the inputs.
@@ -215,8 +219,8 @@ class StableDiffusion3Backbone(Backbone):
             model. Defaults to `1000`.
         shift: float. The shift value for the timestep schedule. Defaults to
             `3.0`.
-        height: optional int. The output height of the image.
-        width: optional int. The output width of the image.
+        image_shape: tuple. The input shape without the batch size. Defaults to
+            `(1024, 1024, 3)`.
         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
@@ -248,6 +252,7 @@ class StableDiffusion3Backbone(Backbone):
         mmdit_hidden_dim=256,
         mmdit_depth=4,
         mmdit_position_size=192,
+        mmdit_qk_norm=None,
         vae=vae,
         clip_l=clip_l,
         clip_g=clip_g,
@@ -262,6 +267,7 @@ class StableDiffusion3Backbone(Backbone):
         mmdit_num_layers,
         mmdit_num_heads,
         mmdit_position_size,
+        mmdit_qk_norm,
         vae,
         clip_l,
         clip_g,
@@ -270,23 +276,21 @@ class StableDiffusion3Backbone(Backbone):
         output_channels=3,
         num_train_timesteps=1000,
         shift=3.0,
-        height=None,
-        width=None,
+        image_shape=(1024, 1024, 3),
         data_format=None,
         dtype=None,
         **kwargs,
     ):
-        height = int(height or 1024)
-        width = int(width or 1024)
-        if height % 8 != 0 or width % 8 != 0:
-            raise ValueError(
-                "`height` and `width` must be divisible by 8. "
-                f"Received: height={height}, width={width}"
-            )
         data_format = standardize_data_format(data_format)
         if data_format != "channels_last":
             raise NotImplementedError
-        image_shape = (height, width, int(vae.input_channels))
+        height = image_shape[0]
+        width = image_shape[1]
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(
+                "height and width in `image_shape` must be divisible by 8. "
+                f"Received: image_shape={image_shape}"
+            )
         latent_shape = (height // 8, width // 8, int(latent_channels))
         context_shape = (None, 4096 if t5 is None else t5.hidden_dim)
         pooled_projection_shape = (clip_l.hidden_dim + clip_g.hidden_dim,)
@@ -314,6 +318,7 @@ class StableDiffusion3Backbone(Backbone):
             latent_shape=latent_shape,
             context_shape=context_shape,
             pooled_projection_shape=pooled_projection_shape,
+            qk_norm=mmdit_qk_norm,
             data_format=data_format,
             dtype=dtype,
             name="diffuser",
@@ -448,12 +453,12 @@ class StableDiffusion3Backbone(Backbone):
         self.mmdit_num_layers = mmdit_num_layers
         self.mmdit_num_heads = mmdit_num_heads
         self.mmdit_position_size = mmdit_position_size
+        self.mmdit_qk_norm = mmdit_qk_norm
         self.latent_channels = latent_channels
         self.output_channels = output_channels
         self.num_train_timesteps = num_train_timesteps
         self.shift = shift
-        self.height = height
-        self.width = width
+        self.image_shape = image_shape
 
     @property
     def latent_shape(self):
@@ -535,7 +540,7 @@ class StableDiffusion3Backbone(Backbone):
         embeddings,
         step,
         num_steps,
-        guidance_scale,
+        guidance_scale=None,
     ):
         step = ops.convert_to_tensor(step)
         next_step = ops.add(step, 1)
@@ -543,9 +548,15 @@ class StableDiffusion3Backbone(Backbone):
         next_sigma, _ = self.scheduler(next_step, num_steps)
 
         # Concatenation for classifier-free guidance.
-        concated_latents, contexts, pooled_projs, timesteps = self.cfg_concat(
-            latents, *embeddings, timestep
-        )
+        if guidance_scale is not None:
+            concated_latents, contexts, pooled_projs, timesteps = (
+                self.cfg_concat(latents, *embeddings, timestep)
+            )
+        else:
+            timesteps = ops.broadcast_to(timestep, ops.shape(latents)[:1])
+            concated_latents = latents
+            contexts = embeddings[0]
+            pooled_projs = embeddings[2]
 
         # Diffusion.
         predicted_noise = self.diffuser(
@@ -559,7 +570,8 @@ class StableDiffusion3Backbone(Backbone):
         )
 
         # Classifier-free guidance.
-        predicted_noise = self.cfg(predicted_noise, guidance_scale)
+        if guidance_scale is not None:
+            predicted_noise = self.cfg(predicted_noise, guidance_scale)
 
         # Euler step.
         return self.euler_step(latents, predicted_noise, sigma, next_sigma)
@@ -577,6 +589,7 @@ class StableDiffusion3Backbone(Backbone):
                 "mmdit_num_layers": self.mmdit_num_layers,
                 "mmdit_num_heads": self.mmdit_num_heads,
                 "mmdit_position_size": self.mmdit_position_size,
+                "mmdit_qk_norm": self.mmdit_qk_norm,
                 "vae": layers.serialize(self.vae),
                 "clip_l": layers.serialize(self.clip_l),
                 "clip_g": layers.serialize(self.clip_g),
@@ -585,8 +598,7 @@ class StableDiffusion3Backbone(Backbone):
                 "output_channels": self.output_channels,
                 "num_train_timesteps": self.num_train_timesteps,
                 "shift": self.shift,
-                "height": self.height,
-                "width": self.width,
+                "image_shape": self.image_shape,
             }
         )
         return config
@@ -624,4 +636,9 @@ class StableDiffusion3Backbone(Backbone):
             config["t5"] = layers.deserialize(
                 config["t5"], custom_objects=custom_objects
             )
+
+        # To maintain backward compatibility, we need to ensure that
+        # `mmdit_qk_norm` is included in the config.
+        if "mmdit_qk_norm" not in config:
+            config["mmdit_qk_norm"] = None
         return cls(**config)

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_image_to_image.py

@@ -27,7 +27,7 @@ class StableDiffusion3ImageToImage(ImageToImage):
     Use `generate()` to do image generation.
     ```python
     image_to_image = keras_hub.models.StableDiffusion3ImageToImage.from_preset(
-        "stable_diffusion_3_medium", height=512, width=512
+        "stable_diffusion_3_medium", image_shape=(512, 512, 3)
     )
     image_to_image.generate(
         {
@@ -158,14 +158,14 @@ class StableDiffusion3ImageToImage(ImageToImage):
         self,
         inputs,
         num_steps=50,
-        guidance_scale=7.0,
         strength=0.8,
+        guidance_scale=7.0,
         seed=None,
     ):
         return super().generate(
             inputs,
             num_steps=num_steps,
-            guidance_scale=guidance_scale,
             strength=strength,
+            guidance_scale=guidance_scale,
             seed=seed,
         )

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_inpaint.py

@@ -29,7 +29,7 @@ class StableDiffusion3Inpaint(Inpaint):
     reference_image = np.ones((1024, 1024, 3), dtype="float32")
     reference_mask = np.ones((1024, 1024), dtype="float32")
     inpaint = keras_hub.models.StableDiffusion3Inpaint.from_preset(
-        "stable_diffusion_3_medium", height=512, width=512
+        "stable_diffusion_3_medium", image_shape=(512, 512, 3)
     )
     inpaint.generate(
         reference_image,
@@ -181,14 +181,14 @@ class StableDiffusion3Inpaint(Inpaint):
         self,
         inputs,
         num_steps=50,
-        guidance_scale=7.0,
         strength=0.6,
+        guidance_scale=7.0,
         seed=None,
     ):
         return super().generate(
             inputs,
             num_steps=num_steps,
-            guidance_scale=guidance_scale,
             strength=strength,
+            guidance_scale=guidance_scale,
             seed=seed,
         )

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_presets.py

@@ -9,10 +9,34 @@ backbone_presets = {
                 "Developed by Stability AI."
             ),
             "params": 2987080931,
-            "official_name": "StableDiffusion3",
             "path": "stable_diffusion_3",
-            "model_card": "https://arxiv.org/abs/2110.00476",
         },
-        "kaggle_handle": "kaggle://keras/stablediffusion3/keras/stable_diffusion_3_medium/1",
-    }
+        "kaggle_handle": "kaggle://keras/stablediffusion3/keras/stable_diffusion_3_medium/3",
+    },
+    "stable_diffusion_3.5_large": {
+        "metadata": {
+            "description": (
+                "9 billion parameter, including CLIP L and CLIP G text "
+                "encoders, MMDiT generative model, and VAE autoencoder. "
+                "Developed by Stability AI."
+            ),
+            "params": 9048410595,
+            "path": "stable_diffusion_3",
+        },
+        "kaggle_handle": "kaggle://keras/stablediffusion-3.5/keras/stable_diffusion_3.5_large/1",
+    },
+    "stable_diffusion_3.5_large_turbo": {
+        "metadata": {
+            "description": (
+                "9 billion parameter, including CLIP L and CLIP G text "
+                "encoders, MMDiT generative model, and VAE autoencoder. "
+                "A timestep-distilled version that eliminates classifier-free "
+                "guidance and uses fewer steps for generation. "
+                "Developed by Stability AI."
+            ),
+            "params": 9048410595,
+            "path": "stable_diffusion_3",
+        },
+        "kaggle_handle": "kaggle://keras/stablediffusion-3.5/keras/stable_diffusion_3.5_large_turbo/1",
+    },
 }

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image.py

@@ -27,7 +27,7 @@ class StableDiffusion3TextToImage(TextToImage):
     Use `generate()` to do image generation.
     ```python
     text_to_image = keras_hub.models.StableDiffusion3TextToImage.from_preset(
-        "stable_diffusion_3_medium", height=512, width=512
+        "stable_diffusion_3_medium", image_shape=(512, 512, 3)
     )
     text_to_image.generate(
         "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

keras_hub/src/models/t5/t5_backbone.py

@@ -42,11 +42,12 @@ class T5Backbone(Backbone):
             projections in the multi-head attention layers. Defaults to
             hidden_dim / num_heads.
         dropout: float. Dropout probability for the Transformer layers.
-        activation: activation function (or activation string name). The
-            activation to be used in the inner dense blocks of the
-            Transformer layers. Defaults to `"relu"`.
+        activation: string. The activation function to use in the dense blocks
+            of the Transformer Layers.
         use_gated_activation: boolean. Whether to use activation gating in
-            the inner dense blocks of the Transformer layers.
+            the inner dense blocks of the Transformer layers. When used with
+            the GELU activation function, this is referred to as GEGLU
+            (gated GLU) from https://arxiv.org/pdf/2002.05202.
             The original T5 architecture didn't use gating, but more
             recent versions do. Defaults to `True`.
         layer_norm_epsilon: float. Epsilon factor to be used in the

keras_hub/src/models/t5/t5_presets.py

@@ -1,4 +1,4 @@
-"""XLM-RoBERTa model preset configurations."""
+"""T5 model preset configurations."""
 
 backbone_presets = {
     "t5_small_multi": {
@@ -8,12 +8,18 @@ backbone_presets = {
                 "Corpus (C4)."
             ),
             "params": 0,
-            "official_name": "T5",
             "path": "t5",
-            "model_card": "https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.md",
         },
         "kaggle_handle": "kaggle://keras/t5/keras/t5_small_multi/2",
     },
+    "t5_1.1_small": {
+        "metadata": {
+            "description": (""),
+            "params": 60511616,
+            "path": "t5",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_small/1",
+    },
     "t5_base_multi": {
         "metadata": {
             "description": (
@@ -21,12 +27,18 @@ backbone_presets = {
                 "Corpus (C4)."
             ),
             "params": 0,
-            "official_name": "T5",
             "path": "t5",
-            "model_card": "https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.md",
         },
         "kaggle_handle": "kaggle://keras/t5/keras/t5_base_multi/2",
     },
+    "t5_1.1_base": {
+        "metadata": {
+            "description": (""),
+            "params": 247577856,
+            "path": "t5",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_base/1",
+    },
     "t5_large_multi": {
         "metadata": {
             "description": (
@@ -34,12 +46,34 @@ backbone_presets = {
                 "Corpus (C4)."
             ),
             "params": 0,
-            "official_name": "T5",
             "path": "t5",
-            "model_card": "https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.md",
         },
         "kaggle_handle": "kaggle://keras/t5/keras/t5_large_multi/2",
     },
+    "t5_1.1_large": {
+        "metadata": {
+            "description": (""),
+            "params": 750251008,
+            "path": "t5",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_large/1",
+    },
+    "t5_1.1_xl": {
+        "metadata": {
+            "description": (""),
+            "params": 2849757184,
+            "path": "t5",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_xl/1",
+    },
+    "t5_1.1_xxl": {
+        "metadata": {
+            "description": (""),
+            "params": 11135332352,
+            "path": "t5",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_xxl/1",
+    },
     "flan_small_multi": {
         "metadata": {
             "description": (
@@ -47,9 +81,7 @@ backbone_presets = {
                 "Corpus (C4)."
             ),
             "params": 0,
-            "official_name": "T5",
             "path": "t5",
-            "model_card": "https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.md",
         },
         "kaggle_handle": "kaggle://keras/t5/keras/flan_small_multi/2",
     },
@@ -60,9 +92,7 @@ backbone_presets = {
                 "Corpus (C4)."
             ),
             "params": 0,
-            "official_name": "T5",
             "path": "t5",
-            "model_card": "https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.md",
         },
         "kaggle_handle": "kaggle://keras/t5/keras/flan_base_multi/2",
     },
@@ -73,9 +103,7 @@ backbone_presets = {
                 "Corpus (C4)."
             ),
             "params": 0,
-            "official_name": "T5",
             "path": "t5",
-            "model_card": "https://github.com/google-research/text-to-text-transfer-transformer/blob/main/README.md",
         },
         "kaggle_handle": "kaggle://keras/t5/keras/flan_large_multi/2",
     },

keras_hub/src/models/text_to_image.py

@@ -249,7 +249,7 @@ class TextToImage(Task):
         self,
         inputs,
         num_steps,
-        guidance_scale,
+        guidance_scale=None,
         seed=None,
     ):
         """Generate image based on the provided `inputs`.
@@ -283,15 +283,23 @@ class TextToImage(Task):
                 - A `tf.data.Dataset` with "prompts" and/or "negative_prompts"
                     keys
             num_steps: int. The number of diffusion steps to take.
-            guidance_scale: float. The classifier free guidance scale defined in
-                [Classifier-Free Diffusion Guidance](
+            guidance_scale: Optional float. The classifier free guidance scale
+                defined in [Classifier-Free Diffusion Guidance](
                 https://arxiv.org/abs/2207.12598). A higher scale encourages
                 generating images more closely related to the prompts, typically
-                at the cost of lower image quality.
+                at the cost of lower image quality. Note that some models don't
+                utilize classifier-free guidance.
             seed: optional int. Used as a random seed.
         """
+        num_steps = int(num_steps)
+        guidance_scale = (
+            float(guidance_scale) if guidance_scale is not None else None
+        )
         num_steps = ops.convert_to_tensor(num_steps, "int32")
-        guidance_scale = ops.convert_to_tensor(guidance_scale)
+        if guidance_scale is not None and guidance_scale > 1.0:
+            guidance_scale = ops.convert_to_tensor(guidance_scale)
+        else:
+            guidance_scale = None
 
         # Setup our three main passes.
         # 1. Preprocessing strings to dense integer tensors.

keras_hub/src/models/vgg/vgg_backbone.py

@@ -27,7 +27,7 @@ class VGGBackbone(Backbone):
     input_data = np.ones((2, 224, 224, 3), dtype="float32")
 
     # Pretrained VGG backbone.
-    model = keras_hub.models.VGGBackbone.from_preset("vgg16")
+    model = keras_hub.models.VGGBackbone.from_preset("vgg_16_imagenet")
     model(input_data)
 
     # Randomly initialized VGG backbone with a custom config.

keras_hub/src/models/vgg/vgg_presets.py

@@ -8,9 +8,7 @@ backbone_presets = {
                 "at a 224x224 resolution."
             ),
             "params": 9220480,
-            "official_name": "vgg",
             "path": "vgg",
-            "model_card": "https://arxiv.org/abs/1409.1556",
         },
         "kaggle_handle": "kaggle://keras/vgg/keras/vgg_11_imagenet/1",
     },
@@ -21,9 +19,7 @@ backbone_presets = {
                 "at a 224x224 resolution."
             ),
             "params": 9404992,
-            "official_name": "vgg",
             "path": "vgg",
-            "model_card": "https://arxiv.org/abs/1409.1556",
         },
         "kaggle_handle": "kaggle://keras/vgg/keras/vgg_13_imagenet/1",
     },
@@ -34,9 +30,7 @@ backbone_presets = {
                 "at a 224x224 resolution."
             ),
             "params": 14714688,
-            "official_name": "vgg",
             "path": "vgg",
-            "model_card": "https://arxiv.org/abs/1409.1556",
         },
         "kaggle_handle": "kaggle://keras/vgg/keras/vgg_16_imagenet/1",
     },
@@ -47,9 +41,7 @@ backbone_presets = {
                 "at a 224x224 resolution."
             ),
             "params": 20024384,
-            "official_name": "vgg",
             "path": "vgg",
-            "model_card": "https://arxiv.org/abs/1409.1556",
         },
         "kaggle_handle": "kaggle://keras/vgg/keras/vgg_19_imagenet/1",
     },

keras_hub/src/models/whisper/whisper_audio_converter.py

@@ -39,7 +39,7 @@ class WhisperAudioConverter(AudioConverter):
     audio_tensor = tf.ones((8000,), dtype="float32")
 
     # Compute the log-mel spectrogram.
-    audio_converter = keras_hub.models.WhisperAudioConverter.from_preset(
+    audio_converter = keras_hub.layers.WhisperAudioConverter.from_preset(
         "whisper_base_en",
     )
     audio_converter(audio_tensor)