Merge pull request #41 from argmaxinc/sd3.5_large

Model support: stable-diffusion-3.5-large

Author: arda-argmax

python/src/diffusionkit/mlx/__init__.py

@@ -36,14 +36,15 @@
 
 MMDIT_CKPT = {
     "argmaxinc/mlx-stable-diffusion-3-medium": "argmaxinc/mlx-stable-diffusion-3-medium",
-    "sd3-8b-unreleased": "models/sd3_8b_beta.safetensors",  # unreleased
+    "argmaxinc/mlx-stable-diffusion-3.5-large": "argmaxinc/mlx-stable-diffusion-3.5-large",
     "argmaxinc/mlx-FLUX.1-schnell": "argmaxinc/mlx-FLUX.1-schnell",
     "argmaxinc/mlx-FLUX.1-schnell-4bit-quantized": "argmaxinc/mlx-FLUX.1-schnell-4bit-quantized",
     "argmaxinc/mlx-FLUX.1-dev": "argmaxinc/mlx-FLUX.1-dev",
 }
 
 T5_MAX_LENGTH = {
     "argmaxinc/mlx-stable-diffusion-3-medium": 512,
+    "argmaxinc/mlx-stable-diffusion-3.5-large": 512,
     "argmaxinc/mlx-FLUX.1-schnell": 256,
     "argmaxinc/mlx-FLUX.1-schnell-4bit-quantized": 256,
     "argmaxinc/mlx-FLUX.1-dev": 512,

python/src/diffusionkit/mlx/config.py

@@ -71,7 +71,9 @@ def hidden_size(self) -> int:
     guidance_embed: bool = False
 
 
-SD3_8b = MMDiTConfig(depth_multimodal=38, num_heads=3, upcast_multimodal_blocks=[35])
+SD3_8b = MMDiTConfig(
+    depth_multimodal=38, num_heads=38, upcast_multimodal_blocks=[35], use_qk_norm=True
+)
 
 SD3_2b = MMDiTConfig(
     depth_multimodal=24, num_heads=24, float16_dtype=mx.float16, dtype=mx.float16

python/src/diffusionkit/mlx/mmdit.py

@@ -818,7 +818,7 @@ def __init__(
         self.kv_proj_embed_dim = self.per_head_dim * n_heads
 
         # Note: key bias is redundant due to softmax invariance
-        self.k_proj = nn.Linear(embed_dim, self.kv_proj_embed_dim)
+        self.k_proj = nn.Linear(embed_dim, self.kv_proj_embed_dim, bias=False)
         self.q_proj = nn.Linear(embed_dim, embed_dim)
         self.v_proj = nn.Linear(embed_dim, self.kv_proj_embed_dim)
         self.o_proj = nn.Linear(embed_dim, embed_dim)

python/src/diffusionkit/mlx/model_io.py

@@ -20,6 +20,7 @@
     AutoencoderConfig,
     CLIPTextModelConfig,
     SD3_2b,
+    SD3_8b,
     VAEDecoderConfig,
     VAEEncoderConfig,
 )
@@ -50,6 +51,10 @@
         "argmaxinc/mlx-FLUX.1-dev": "flux1-dev.safetensors",
         "vae": "ae.safetensors",
     },
+    "argmaxinc/mlx-stable-diffusion-3.5-large": {
+        "argmaxinc/mlx-stable-diffusion-3.5-large": "sd3.5_large.safetensors",
+        "vae": "sd3.5_large.safetensors",
+    },
 }
 _DEFAULT_MODEL = "argmaxinc/stable-diffusion"
 _MODELS = {
@@ -83,17 +88,29 @@
         "vae_encoder": "encoder.",
         "vae_decoder": "decoder.",
     },
+    "argmaxinc/mlx-stable-diffusion-3.5-large": {
+        "vae_encoder": "first_stage_model.encoder.",
+        "vae_decoder": "first_stage_model.decoder.",
+    },
+}
+
+_CONFIG = {
+    "argmaxinc/mlx-stable-diffusion-3-medium": SD3_2b,
+    "argmaxinc/mlx-FLUX.1-schnell": FLUX_SCHNELL,
+    "argmaxinc/mlx-FLUX.1-schnell-4bit-quantized": FLUX_SCHNELL,
+    "argmaxinc/mlx-FLUX.1-dev": FLUX_SCHNELL,
+    "argmaxinc/mlx-stable-diffusion-3.5-large": SD3_8b,
 }
 
 _FLOAT16 = mx.bfloat16
 
 DEPTH = {
     "argmaxinc/mlx-stable-diffusion-3-medium": 24,
-    "sd3-8b-unreleased": 38,
+    "argmaxinc/mlx-stable-diffusion-3.5-large": 38,
 }
 MAX_LATENT_RESOLUTION = {
     "argmaxinc/mlx-stable-diffusion-3-medium": 96,
-    "sd3-8b-unreleased": 192,
+    "argmaxinc/mlx-stable-diffusion-3.5-large": 192,
 }
 
 LOCAl_SD3_CKPT = None
@@ -321,6 +338,14 @@ def mmdit_state_dict_adjustments(state_dict, prefix=""):
         for k, v in state_dict.items()
     }
 
+    # Remap qk_norm
+    state_dict = {
+        k.replace(".attn.ln_q.", ".qk_norm.q_norm."): v for k, v in state_dict.items()
+    }
+    state_dict = {
+        k.replace(".attn.ln_k.", ".qk_norm.k_norm."): v for k, v in state_dict.items()
+    }
+
     # Split qkv proj and rename:
     # *transformer_block.attn.qkv.{weigth/bias}  -> transformer_block.attn.{q/k/v}_proj.{weigth/bias}
     # *transformer_block.attn.proj.{weigth/bias} -> transformer_block.attn.o_proj.{weight/bias}
@@ -347,6 +372,9 @@ def mmdit_state_dict_adjustments(state_dict, prefix=""):
     # Filter out VAE Decoder related tensors
     state_dict = {k: v for k, v in state_dict.items() if "decoder." not in k}
 
+    # Filter out VAE Encoder related tensors
+    state_dict = {k: v for k, v in state_dict.items() if "encoder." not in k}
+
     # Filter out k_proj.bias related tensors
     state_dict = {k: v for k, v in state_dict.items() if "k_proj.bias" not in k}
 
@@ -676,7 +704,7 @@ def load_mmdit(
     """Load the MM-DiT model from the checkpoint file."""
     """only_modulation_dict: Only returns the modulation dictionary"""
     dtype = _FLOAT16 if float16 else mx.float32
-    config = SD3_2b
+    config = _CONFIG[key]
     config.low_memory_mode = low_memory_mode
     model = MMDiT(config)
 

python/src/diffusionkit/mlx/scripts/generate_images.py

@@ -14,21 +14,21 @@
 # Defaults
 HEIGHT = {
     "argmaxinc/mlx-stable-diffusion-3-medium": 512,
-    "sd3-8b-unreleased": 1024,
+    "argmaxinc/mlx-stable-diffusion-3.5-large": 1024,
     "argmaxinc/mlx-FLUX.1-schnell": 512,
     "argmaxinc/mlx-FLUX.1-schnell-4bit-quantized": 512,
     "argmaxinc/mlx-FLUX.1-dev": 512,
 }
 WIDTH = {
     "argmaxinc/mlx-stable-diffusion-3-medium": 512,
-    "sd3-8b-unreleased": 1024,
+    "argmaxinc/mlx-stable-diffusion-3.5-large": 1024,
     "argmaxinc/mlx-FLUX.1-schnell": 512,
     "argmaxinc/mlx-FLUX.1-schnell-4bit-quantized": 512,
     "argmaxinc/mlx-FLUX.1-dev": 512,
 }
 SHIFT = {
     "argmaxinc/mlx-stable-diffusion-3-medium": 3.0,
-    "sd3-8b-unreleased": 3.0,
+    "argmaxinc/mlx-stable-diffusion-3.5-large": 3.0,
     "argmaxinc/mlx-FLUX.1-schnell": 1.0,
     "argmaxinc/mlx-FLUX.1-schnell-4bit-quantized": 1.0,
     "argmaxinc/mlx-FLUX.1-dev": 1.0,

setup.py

@@ -3,7 +3,7 @@
 from setuptools import find_packages, setup
 from setuptools.command.install import install
 
-VERSION = "0.4.0"
+VERSION = "0.5.0"
 
 
 class VersionInstallCommand(install):