Add stable diffusion 3.5 medium model (#2033)

Author: james77777778

keras_hub/src/models/stable_diffusion_3/mmdit.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 @@ def __init__(
         mmdit_num_heads,
         mmdit_position_size,
         mmdit_qk_norm,
+        mmdit_dual_attention_indices,
         vae,
         clip_l,
         clip_g,
@@ -319,6 +324,7 @@ def __init__(
             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 @@ def __init__(
         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 @@ def get_config(self):
                 "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 @@ def from_config(cls, config, custom_objects=None):
             )
 
         # 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/src/models/stable_diffusion_3/stable_diffusion_3_backbone.py

@@ -35,6 +35,7 @@ def setUp(self):
             "mmdit_num_heads": 2,
             "mmdit_position_size": 192,
             "mmdit_qk_norm": None,
+            "mmdit_dual_attention_indices": None,
             "vae": vae,
             "clip_l": clip_l,
             "clip_g": clip_g,
@@ -67,10 +68,15 @@ def test_backbone_basics(self):
             run_quantization_check=False,
         )
 
-        # Test `mmdit_qk_norm="rms_norm"`.
+    def test_backbone_basics_mmditx(self):
+        # MMDiT-X includes `mmdit_qk_norm` and `mmdit_dual_attention_indices`.
         self.run_backbone_test(
             cls=StableDiffusion3Backbone,
-            init_kwargs={**self.init_kwargs, "mmdit_qk_norm": "rms_norm"},
+            init_kwargs={
+                **self.init_kwargs,
+                "mmdit_qk_norm": "rms_norm",
+                "mmdit_dual_attention_indices": (0,),
+            },
             input_data=self.input_data,
             expected_output_shape={
                 "images": (2, 64, 64, 3),

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_backbone_test.py

@@ -41,6 +41,7 @@ def setUp(self):
             mmdit_num_heads=2,
             mmdit_position_size=192,
             mmdit_qk_norm=None,
+            mmdit_dual_attention_indices=None,
             vae=VAEBackbone(
                 [32, 32, 32, 32],
                 [1, 1, 1, 1],

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_image_to_image_test.py

@@ -41,6 +41,7 @@ def setUp(self):
             mmdit_num_heads=2,
             mmdit_position_size=192,
             mmdit_qk_norm=None,
+            mmdit_dual_attention_indices=None,
             vae=VAEBackbone(
                 [32, 32, 32, 32],
                 [1, 1, 1, 1],

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_inpaint_test.py

@@ -13,6 +13,18 @@
         },
         "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/src/models/stable_diffusion_3/stable_diffusion_3_presets.py

@@ -41,6 +41,7 @@ def setUp(self):
             mmdit_num_heads=2,
             mmdit_position_size=192,
             mmdit_qk_norm=None,
+            mmdit_dual_attention_indices=None,
             vae=VAEBackbone(
                 [32, 32, 32, 32],
                 [1, 1, 1, 1],

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image_test.py

@@ -6,6 +6,10 @@
 python tools/checkpoint_conversion/convert_stable_diffusion_3_checkpoints.py \
     --preset stable_diffusion_3_medium \
     --upload_uri kaggle://kerashub/stablediffusion3/keras/stable_diffusion_3_medium
+python tools/checkpoint_conversion/convert_stable_diffusion_3_checkpoints.py \
+    --preset stable_diffusion_3.5_medium \
+    --upload_uri kaggle://kerashub/stablediffusion3/keras/stable_diffusion_3.5_medium \
+    --dtype bfloat16
 python tools/checkpoint_conversion/convert_stable_diffusion_3_checkpoints.py \
     --preset stable_diffusion_3.5_large \
     --upload_uri kaggle://kerashub/stablediffusion3/keras/stable_diffusion_3.5_large \
@@ -56,6 +60,17 @@
         # Tokenizer
         "clip_tokenizer": "hf://openai/clip-vit-large-patch14",
     },
+    "stable_diffusion_3.5_medium": {
+        # HF root
+        "root": "hf://stabilityai/stable-diffusion-3.5-medium",
+        # Model <-> Path
+        "clip_l": "text_encoder/model.safetensors",
+        "clip_g": "text_encoder_2/model.safetensors",
+        "diffuser": "sd3.5_medium.safetensors",
+        "vae": "sd3.5_medium.safetensors",
+        # Tokenizer
+        "clip_tokenizer": "hf://openai/clip-vit-large-patch14",
+    },
     "stable_diffusion_3.5_large": {
         # HF root
         "root": "hf://stabilityai/stable-diffusion-3.5-large",
@@ -148,11 +163,27 @@ def convert_model(preset, height, width):
             24,
             192,
             None,  # qk_norm
+            None,  # dual_attention_indices
+            vae,
+            clip_l,
+            clip_g,
+            image_shape=(height, width, 3),
+            name="stable_diffusion_3_medium_backbone",
+        )
+    elif preset == "stable_diffusion_3.5_medium":
+        backbone = StableDiffusion3Backbone(
+            2,
+            64 * 24,
+            24,
+            24,
+            384,  # position_size is larger than SD3
+            "rms_norm",  # qk_norm
+            (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12),  # dual_attn_indices
             vae,
             clip_l,
             clip_g,
             image_shape=(height, width, 3),
-            name="stable_diffusion_3_backbone",
+            name="stable_diffusion_3.5_medium_backbone",
         )
     elif preset in (
         "stable_diffusion_3.5_large",
@@ -165,11 +196,12 @@ def convert_model(preset, height, width):
             38,
             192,
             "rms_norm",  # qk_norm
+            None,  # dual_attention_indices
             vae,
             clip_l,
             clip_g,
             image_shape=(height, width, 3),
-            name="stable_diffusion_3.5_backbone",
+            name="stable_diffusion_3.5_large_backbone",
         )
     else:
         raise ValueError(f"Unknown preset={preset}.")
@@ -418,6 +450,24 @@ def port_diffuser(preset, filename, model):
                     port_dense(loader, block.mlp.dense1, f"{prefix}.mlp.fc1")
                     port_dense(loader, block.mlp.dense2, f"{prefix}.mlp.fc2")
 
+                    # Dual attention
+                    if block.use_dual_attention:
+                        port_dense(
+                            loader, block.attention_qkv2, f"{prefix}.attn2.qkv"
+                        )
+                        if block.qk_norm is not None:
+                            port_ln_or_gn(
+                                loader, block.q_norm2, f"{prefix}.attn2.ln_q"
+                            )
+                            port_ln_or_gn(
+                                loader, block.k_norm2, f"{prefix}.attn2.ln_k"
+                            )
+                        port_dense(
+                            loader,
+                            block.attention_proj2,
+                            f"{prefix}.attn2.proj",
+                        )
+
             # Output layer
             port_dense(
                 loader,
@@ -562,7 +612,10 @@ def validate_output(preset, keras_model, keras_preprocessor, output_dir):
     if preset == "stable_diffusion_3_medium":
         num_steps = 28
         guidance_scale = 7.0
-    elif preset == "stable_diffusion_3.5_large":
+    elif preset in (
+        "stable_diffusion_3.5_medium",
+        "stable_diffusion_3.5_large",
+    ):
         num_steps = 40
         guidance_scale = 4.5
     elif preset == "stable_diffusion_3.5_large_turbo":