keras-hub-nightly 0.16.1.dev202410020340__py3-none-any.whl → 0.19.0.dev202501260345__py3-none-any.whl

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_image_to_image.py

@@ -0,0 +1,178 @@
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.image_to_image import ImageToImage
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (  # noqa: E501
+    StableDiffusion3Backbone,
+)
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (  # noqa: E501
+    StableDiffusion3TextToImagePreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.StableDiffusion3ImageToImage")
+class StableDiffusion3ImageToImage(ImageToImage):
+    """An end-to-end Stable Diffusion 3 model for image-to-image generation.
+
+    This model has a `generate()` method, which generates images based
+    on a combination of a reference image and a text prompt.
+
+    Args:
+        backbone: A `keras_hub.models.StableDiffusion3Backbone` instance.
+        preprocessor: A
+            `keras_hub.models.StableDiffusion3TextToImagePreprocessor` instance.
+
+    Examples:
+
+    Use `generate()` to do image generation.
+    ```python
+    prompt = (
+        "Astronaut in a jungle, cold color palette, muted colors, "
+        "detailed, 8k"
+    )
+    image_to_image = keras_hub.models.StableDiffusion3ImageToImage.from_preset(
+        "stable_diffusion_3_medium", image_shape=(512, 512, 3)
+    )
+    image_to_image.generate(
+        {
+            "images": np.ones((512, 512, 3), dtype="float32"),
+            "prompts": prompt,
+        }
+    )
+
+    # Generate with batched prompts.
+    image_to_image.generate(
+        {
+            "images": np.ones((2, 512, 512, 3), dtype="float32"),
+            "prompts": [
+                "cute wallpaper art of a cat",
+                "cute wallpaper art of a dog",
+            ],
+        }
+    )
+
+    # Generate with different `num_steps`, `guidance_scale` and `strength`.
+    image_to_image.generate(
+        {
+            "images": np.ones((512, 512, 3), dtype="float32"),
+            "prompts": prompt,
+        }
+        num_steps=50,
+        guidance_scale=5.0,
+        strength=0.6,
+    )
+
+    # Generate with `negative_prompts`.
+    text_to_image.generate(
+        {
+            "images": np.ones((512, 512, 3), dtype="float32"),
+            "prompts": prompt,
+            "negative_prompts": "green color",
+        }
+    )
+    ```
+    """
+
+    backbone_cls = StableDiffusion3Backbone
+    preprocessor_cls = StableDiffusion3TextToImagePreprocessor
+
+    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 "
+            "`StableDiffusion3ImageToImage`."
+        )
+
+    def generate_step(
+        self,
+        images,
+        noises,
+        token_ids,
+        starting_step,
+        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:
+            images: A (batch_size, image_height, image_width, 3) tensor
+                containing the reference images.
+            noises: A (batch_size, latent_height, latent_width, channels) tensor
+                containing the noises to be added to the latents. 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.
+            starting_step: int. The number of the starting diffusion step.
+            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 images.
+        latents = self.backbone.encode_image_step(images)
+
+        # Add noises to latents.
+        latents = self.backbone.add_noise_step(
+            latents, noises, starting_step, num_steps
+        )
+
+        # Encode inputs.
+        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(starting_step, num_steps, body_fun, latents)
+
+        # Decode.
+        return self.backbone.decode_step(latents)
+
+    def generate(
+        self,
+        inputs,
+        num_steps=50,
+        strength=0.8,
+        guidance_scale=7.0,
+        seed=None,
+    ):
+        return super().generate(
+            inputs,
+            num_steps=num_steps,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            seed=seed,
+        )

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_inpaint.py

@@ -0,0 +1,193 @@
+from keras import ops
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.inpaint import Inpaint
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (  # noqa: E501
+    StableDiffusion3Backbone,
+)
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (  # noqa: E501
+    StableDiffusion3TextToImagePreprocessor,
+)
+
+
+@keras_hub_export("keras_hub.models.StableDiffusion3Inpaint")
+class StableDiffusion3Inpaint(Inpaint):
+    """An end-to-end Stable Diffusion 3 model for inpaint generation.
+
+    This model has a `generate()` method, which generates images based
+    on a combination of a reference image, mask and a text prompt.
+
+    Args:
+        backbone: A `keras_hub.models.StableDiffusion3Backbone` instance.
+        preprocessor: A
+            `keras_hub.models.StableDiffusion3TextToImagePreprocessor` instance.
+
+    Examples:
+
+    Use `generate()` to do image generation.
+    ```python
+    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", image_shape=(512, 512, 3)
+    )
+    inpaint.generate(
+        reference_image,
+        reference_mask,
+        "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+    )
+
+    # Generate with batched prompts.
+    reference_images = np.ones((2, 512, 512, 3), dtype="float32")
+    reference_mask = np.ones((2, 1024, 1024), dtype="float32")
+    inpaint.generate(
+        reference_images,
+        reference_mask,
+        ["cute wallpaper art of a cat", "cute wallpaper art of a dog"]
+    )
+
+    # Generate with different `num_steps`, `guidance_scale` and `strength`.
+    inpaint.generate(
+        reference_image,
+        reference_mask,
+        "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+        num_steps=50,
+        guidance_scale=5.0,
+        strength=0.6,
+    )
+    ```
+    """
+
+    backbone_cls = StableDiffusion3Backbone
+    preprocessor_cls = StableDiffusion3TextToImagePreprocessor
+
+    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 `StableDiffusion3Inpaint`."
+        )
+
+    def generate_step(
+        self,
+        images,
+        masks,
+        noises,
+        token_ids,
+        starting_step,
+        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:
+            images: A (batch_size, image_height, image_width, 3) tensor
+                containing the reference images.
+            masks: A (batch_size, image_height, image_width) tensor
+                containing the reference masks.
+            noises: A (batch_size, latent_height, latent_width, channels) tensor
+                containing the noises to be added to the latents. 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.
+            starting_step: int. The number of the starting diffusion step.
+            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
+
+        # Get masked images.
+        masks = ops.cast(ops.expand_dims(masks, axis=-1) > 0.5, images.dtype)
+        masks_latent_size = ops.image.resize(
+            masks,
+            (self.backbone.latent_shape[1], self.backbone.latent_shape[2]),
+            interpolation="nearest",
+        )
+
+        # Encode images.
+        image_latents = self.backbone.encode_image_step(images)
+
+        # Add noises to latents.
+        latents = self.backbone.add_noise_step(
+            image_latents, noises, starting_step, num_steps
+        )
+
+        # Encode inputs.
+        embeddings = self.backbone.encode_text_step(
+            token_ids, negative_token_ids
+        )
+
+        # Denoise.
+        def body_fun(step, latents):
+            latents = self.backbone.denoise_step(
+                latents,
+                embeddings,
+                step,
+                num_steps,
+                guidance_scale,
+            )
+
+            # Compute the previous latents x_t -> x_t-1.
+            def true_fn():
+                next_step = ops.add(step, 1)
+                return self.backbone.add_noise_step(
+                    image_latents, noises, next_step, num_steps
+                )
+
+            init_latents = ops.cond(
+                step < ops.subtract(num_steps, 1),
+                true_fn,
+                lambda: ops.cast(image_latents, noises.dtype),
+            )
+            latents = ops.add(
+                ops.multiply(
+                    ops.subtract(1.0, masks_latent_size), init_latents
+                ),
+                ops.multiply(masks_latent_size, latents),
+            )
+            return latents
+
+        latents = ops.fori_loop(starting_step, num_steps, body_fun, latents)
+
+        # Decode.
+        return self.backbone.decode_step(latents)
+
+    def generate(
+        self,
+        inputs,
+        num_steps=50,
+        strength=0.6,
+        guidance_scale=7.0,
+        seed=None,
+    ):
+        return super().generate(
+            inputs,
+            num_steps=num_steps,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            seed=seed,
+        )

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_presets.py

@@ -5,14 +5,50 @@ backbone_presets = {
         "metadata": {
             "description": (
                 "3 billion parameter, including CLIP L and CLIP G text "
-                "encoders, MMDiT generative model, and VAE decoder. "
+                "encoders, MMDiT generative model, and VAE autoencoder. "
                 "Developed by Stability AI."
             ),
-            "params": 2952806723,
-            "official_name": "StableDiffusion3",
-            "path": "stablediffusion3",
-            "model_card": "https://arxiv.org/abs/2110.00476",
+            "params": 2987080931,
+            "path": "stable_diffusion_3",
         },
-        "kaggle_handle": "kaggle://kerashub/stablediffusion3/keras/stable_diffusion_3_medium/1",
-    }
+        "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": (
+                "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/2",
+    },
+    "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/2",
+    },
 }

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image.py

@@ -1,10 +1,10 @@
 from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (  # noqa: E501
     StableDiffusion3Backbone,
 )
-from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_text_to_image_preprocessor import (  # noqa: E501
     StableDiffusion3TextToImagePreprocessor,
 )
 from keras_hub.src.models.text_to_image import TextToImage
@@ -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"
@@ -38,11 +38,23 @@ class StableDiffusion3TextToImage(TextToImage):
         ["cute wallpaper art of a cat", "cute wallpaper art of a dog"]
     )
 
-    # Generate with different `num_steps` and `classifier_free_guidance_scale`.
+    # 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,
-        classifier_free_guidance_scale=5.0,
+        guidance_scale=5.0,
+    )
+
+    # Generate with `negative_prompts`.
+    prompt = (
+        "Astronaut in a jungle, cold color palette, muted colors, "
+        "detailed, 8k"
+    )
+    text_to_image.generate(
+        {
+            "prompts": prompt,
+            "negative_prompts": "green color",
+        }
     )
     ```
     """
@@ -79,7 +91,6 @@ class StableDiffusion3TextToImage(TextToImage):
         self,
         latents,
         token_ids,
-        negative_token_ids,
         num_steps,
         guidance_scale,
     ):
@@ -92,10 +103,8 @@ class StableDiffusion3TextToImage(TextToImage):
             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 (batch_size, num_tokens) tensor containing the
-                tokens based on the input prompts.
-            negative_token_ids: A (batch_size, num_tokens) tensor
-                 containing the negative tokens based on the input prompts.
+            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](
@@ -103,8 +112,12 @@ class StableDiffusion3TextToImage(TextToImage):
                 generate images that are closely linked to prompts, usually at
                 the expense of lower image quality.
         """
-        # Encode inputs.
-        embeddings = self.backbone.encode_step(token_ids, negative_token_ids)
+        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):
@@ -124,14 +137,12 @@ class StableDiffusion3TextToImage(TextToImage):
     def generate(
         self,
         inputs,
-        negative_inputs=None,
         num_steps=28,
         guidance_scale=7.0,
         seed=None,
     ):
         return super().generate(
             inputs,
-            negative_inputs=negative_inputs,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image_preprocessor.py

@@ -3,7 +3,7 @@ from keras import layers
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.models.preprocessor import Preprocessor
-from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (
+from keras_hub.src.models.stable_diffusion_3.stable_diffusion_3_backbone import (  # noqa: E501
     StableDiffusion3Backbone,
 )
 

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,11 +8,17 @@ 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",
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_small_multi/3",
+    },
+    "t5_1.1_small": {
+        "metadata": {
+            "description": (""),
+            "params": 60511616,
+            "path": "t5",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_small/2",
     },
     "t5_base_multi": {
         "metadata": {
@@ -21,11 +27,17 @@ 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",
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_base_multi/3",
+    },
+    "t5_1.1_base": {
+        "metadata": {
+            "description": (""),
+            "params": 247577856,
+            "path": "t5",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_base/2",
     },
     "t5_large_multi": {
         "metadata": {
@@ -34,11 +46,33 @@ 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",
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_large_multi/3",
+    },
+    "t5_1.1_large": {
+        "metadata": {
+            "description": (""),
+            "params": 750251008,
+            "path": "t5",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_large/2",
+    },
+    "t5_1.1_xl": {
+        "metadata": {
+            "description": (""),
+            "params": 2849757184,
+            "path": "t5",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_xl/2",
+    },
+    "t5_1.1_xxl": {
+        "metadata": {
+            "description": (""),
+            "params": 11135332352,
+            "path": "t5",
+        },
+        "kaggle_handle": "kaggle://keras/t5/keras/t5_1.1_xxl/2",
     },
     "flan_small_multi": {
         "metadata": {
@@ -47,11 +81,9 @@ 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",
+        "kaggle_handle": "kaggle://keras/t5/keras/flan_small_multi/3",
     },
     "flan_base_multi": {
         "metadata": {
@@ -60,11 +92,9 @@ 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",
+        "kaggle_handle": "kaggle://keras/t5/keras/flan_base_multi/3",
     },
     "flan_large_multi": {
         "metadata": {
@@ -73,10 +103,8 @@ 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",
+        "kaggle_handle": "kaggle://keras/t5/keras/flan_large_multi/3",
     },
 }