Final tests for SD3 and FLUX

Author: atiorh

python/src/diffusionkit/mlx/__init__.py

@@ -495,7 +495,7 @@ def read_image(self, image_path: str):
         # Make sure image shape is divisible by 64
         W, H = (dim - dim % 64 for dim in (img.width, img.height))
         if W != img.width or H != img.height:
-            print(
+            logger.warning(
                 f"Warning: image shape is not divisible by 64, downsampling to {W}x{H}"
             )
             img = img.resize((W, H), Image.LANCZOS)  # use desired downsampling filter
@@ -629,13 +629,21 @@ def __init__(self, model: DiffusionPipeline):
         self.model = model
 
     def cache_modulation_params(self, pooled_text_embeddings, sigmas):
-        self.model.mmdit.cache_modulation_params(pooled_text_embeddings, sigmas)
+        self.model.mmdit.cache_modulation_params(
+            pooled_text_embeddings, sigmas.astype(self.model.activation_dtype)
+        )
 
     def clear_cache(self):
         self.model.mmdit.clear_modulation_params_cache()
 
     def __call__(
-        self, x_t, t, conditioning, cfg_weight: float = 7.5, pooled_conditioning=None
+        self,
+        x_t,
+        timestep,
+        sigma,
+        conditioning,
+        cfg_weight: float = 7.5,
+        pooled_conditioning=None,
     ):
         if cfg_weight <= 0:
             logger.debug("CFG Weight disabled")
@@ -644,20 +652,14 @@ def __call__(
             x_t_mmdit = mx.concatenate([x_t] * 2, axis=0).astype(
                 self.model.activation_dtype
             )
-        t_mmdit = mx.broadcast_to(t, [len(x_t_mmdit)])
-        timestep = self.model.sampler.timestep(t_mmdit).astype(
-            self.model.activation_dtype
-        )
         mmdit_input = {
             "latent_image_embeddings": x_t_mmdit,
             "token_level_text_embeddings": mx.expand_dims(conditioning, 2),
-            "timestep": timestep,
+            "timestep": mx.broadcast_to(timestep, [len(x_t_mmdit)]),
         }
 
         mmdit_output = self.model.mmdit(**mmdit_input)
-        eps_pred = self.model.sampler.calculate_denoised(
-            t_mmdit, mmdit_output, x_t_mmdit
-        )
+        eps_pred = self.model.sampler.calculate_denoised(sigma, mmdit_output, x_t_mmdit)
         if cfg_weight <= 0:
             return eps_pred
         else:
@@ -707,21 +709,22 @@ def to_d(x, sigma, denoised):
 def sample_euler(model: CFGDenoiser, x, sigmas, extra_args=None):
     """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
     extra_args = {} if extra_args is None else extra_args
-    s_in = mx.ones([x.shape[0]])
+
     from tqdm import trange
 
-    sigmas = mx.array([1.0, 0.75, 0.5, 0.25, 0.0], mx.bfloat16)  # FIXME
     t = trange(len(sigmas) - 1)
 
-    model.cache_modulation_params(extra_args.pop("pooled_conditioning"), sigmas)
+    timesteps = model.model.sampler.timestep(sigmas).astype(
+        model.model.activation_dtype
+    )
+    model.cache_modulation_params(extra_args.pop("pooled_conditioning"), timesteps)
 
     iter_time = []
     for i in t:
         start_time = t.format_dict["elapsed"]
-        sigma_hat = sigmas[i]
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        d = to_d(x, sigma_hat, denoised)
-        dt = sigmas[i + 1] - sigma_hat
+        denoised = model(x, timesteps[i], sigmas[i], **extra_args)
+        d = to_d(x, sigmas[i], denoised)
+        dt = sigmas[i + 1] - sigmas[i]
         # Euler method
         x = x + d * dt
         mx.eval(x)

python/src/diffusionkit/mlx/mmdit.py

@@ -3,12 +3,10 @@
 # Copyright (C) 2024 Argmax, Inc. All Rights Reserved.
 #
 
-import gc
 from functools import partial
 
 import mlx.core as mx
 import mlx.nn as nn
-import mlx.utils as utils
 import numpy as np
 from argmaxtools.utils import get_logger
 from beartype.typing import Dict, List, Optional, Tuple
@@ -77,35 +75,31 @@ def cache_modulation_params(
         by offloading all adaLN_modulation parameters
         """
         y_embed = self.y_embedder(pooled_text_embeddings)
+        batch_size = pooled_text_embeddings.shape[0]
 
         offload_size = 0
         to_offload = []
 
         for timestep in timesteps:
             final_timestep = timestep.item() == timesteps[-1].item()
-            modulation_inputs = y_embed + self.t_embedder(timestep[None] * 1000.0)
+            timestep_key = timestep.item()
+            modulation_inputs = y_embed[:, None, None, :] + self.t_embedder(
+                mx.repeat(timestep[None], batch_size, axis=0)
+            )
 
             for block in self.multimodal_transformer_blocks:
                 if not hasattr(block.image_transformer_block, "_modulation_params"):
                     block.image_transformer_block._modulation_params = dict()
                     block.text_transformer_block._modulation_params = dict()
 
                 block.image_transformer_block._modulation_params[
-                    (timestep * 1000).item()
+                    timestep_key
                 ] = block.image_transformer_block.adaLN_modulation(modulation_inputs)
                 block.text_transformer_block._modulation_params[
-                    (timestep * 1000).item()
+                    timestep_key
                 ] = block.text_transformer_block.adaLN_modulation(modulation_inputs)
-                mx.eval(
-                    block.image_transformer_block._modulation_params[
-                        (timestep * 1000).item()
-                    ]
-                )
-                mx.eval(
-                    block.text_transformer_block._modulation_params[
-                        (timestep * 1000).item()
-                    ]
-                )
+                mx.eval(block.image_transformer_block._modulation_params[timestep_key])
+                mx.eval(block.text_transformer_block._modulation_params[timestep_key])
 
                 if final_timestep:
                     offload_size += (
@@ -131,33 +125,34 @@ def cache_modulation_params(
                         ]
                     )
 
-            for block in self.unified_transformer_blocks:
-                if not hasattr(block.transformer_block, "_modulation_params"):
-                    block.transformer_block._modulation_params = dict()
-                block.transformer_block._modulation_params[
-                    (timestep * 1000).item()
-                ] = block.transformer_block.adaLN_modulation(modulation_inputs)
-                mx.eval(
-                    block.transformer_block._modulation_params[(timestep * 1000).item()]
-                )
-
-                if final_timestep:
-                    offload_size += (
-                        block.transformer_block.adaLN_modulation.layers[1].weight.size
-                        * block.transformer_block.adaLN_modulation.layers[
-                            1
-                        ].weight.dtype.size
-                    )
-                    to_offload.extend(
-                        [block.transformer_block.adaLN_modulation.layers[1]]
-                    )
+            if self.config.depth_unified > 0:
+                for block in self.unified_transformer_blocks:
+                    if not hasattr(block.transformer_block, "_modulation_params"):
+                        block.transformer_block._modulation_params = dict()
+                    block.transformer_block._modulation_params[
+                        timestep_key
+                    ] = block.transformer_block.adaLN_modulation(modulation_inputs)
+                    mx.eval(block.transformer_block._modulation_params[timestep_key])
+
+                    if final_timestep:
+                        offload_size += (
+                            block.transformer_block.adaLN_modulation.layers[
+                                1
+                            ].weight.size
+                            * block.transformer_block.adaLN_modulation.layers[
+                                1
+                            ].weight.dtype.size
+                        )
+                        to_offload.extend(
+                            [block.transformer_block.adaLN_modulation.layers[1]]
+                        )
 
             if not hasattr(self.final_layer, "_modulation_params"):
                 self.final_layer._modulation_params = dict()
             self.final_layer._modulation_params[
-                (timestep * 1000).item()
+                timestep_key
             ] = self.final_layer.adaLN_modulation(modulation_inputs)
-            mx.eval(self.final_layer._modulation_params[(timestep * 1000).item()])
+            mx.eval(self.final_layer._modulation_params[timestep_key])
 
             if final_timestep:
                 offload_size += (
@@ -246,6 +241,7 @@ def __call__(
             latent_image_embeddings,
             timestep,
         )
+
         if self.config.patchify_via_reshape:
             latent_image_embeddings = self.x_embedder.unpack(
                 latent_image_embeddings, (latent_height, latent_width)
@@ -437,7 +433,10 @@ def pre_sdpa(
         tensor: mx.array,
         timestep: mx.array,
     ) -> Dict[str, mx.array]:
+        if timestep.size > 1:
+            timestep = timestep[0]
         modulation_params = self._modulation_params[timestep.item()]
+
         modulation_params = mx.split(
             modulation_params, self.num_modulation_params, axis=-1
         )
@@ -771,6 +770,8 @@ def __call__(
         latent_image_embeddings: mx.array,
         timestep: mx.array,
     ) -> mx.array:
+        if timestep.size > 1:
+            timestep = timestep[0]
         modulation_params = self._modulation_params[timestep.item()]
 
         shift, residual_scale = mx.split(modulation_params, 2, axis=-1)