Merge pull request #16 from argmaxinc/benchmark_fix

Benchmark mode fix

Author: atiorh

python/src/diffusionkit/mlx/__init__.py

@@ -67,53 +67,77 @@ def __init__(
         self.dtype = self.float16_dtype if w16 else mx.float32
         self.activation_dtype = self.float16_dtype if a16 else mx.float32
         self.use_t5 = use_t5
-        mmdit_ckpt = MMDIT_CKPT[model_version]
+        self.mmdit_ckpt = MMDIT_CKPT[model_version]
         self.low_memory_mode = low_memory_mode
-        self.mmdit = load_mmdit(
-            float16=w16,
-            key=mmdit_ckpt,
-            model_key=model_version,
-            low_memory_mode=low_memory_mode,
-        )
+        self.model = model
+        self.model_version = model_version
         self.sampler = ModelSamplingDiscreteFlow(shift=shift)
-        self.decoder = load_vae_decoder(float16=w16, key=mmdit_ckpt)
-        self.encoder = load_vae_encoder(float16=False, key=mmdit_ckpt)
         self.latent_format = SD3LatentFormat()
+        self.use_clip_g = True
+        self.check_and_load_models()
 
-        self.clip_l = load_text_encoder(
-            model,
-            w16,
-            model_key="clip_l",
-        )
-        self.tokenizer_l = load_tokenizer(
-            model,
-            merges_key="tokenizer_l_merges",
-            vocab_key="tokenizer_l_vocab",
-            pad_with_eos=True,
-        )
-        self.clip_g = load_text_encoder(
-            model,
-            w16,
-            model_key="clip_g",
-        )
-        self.tokenizer_g = load_tokenizer(
-            model,
-            merges_key="tokenizer_g_merges",
-            vocab_key="tokenizer_g_vocab",
-            pad_with_eos=False,
-        )
-        self.t5_encoder = None
-        self.t5_tokenizer = None
-        if self.use_t5:
+    def load_mmdit(self, only_modulation_dict=False):
+        if only_modulation_dict:
+            return load_mmdit(
+                float16=True if self.dtype == self.float16_dtype else False,
+                key=self.mmdit_ckpt,
+                model_key=self.model_version,
+                low_memory_mode=self.low_memory_mode,
+                only_modulation_dict=only_modulation_dict,
+            )
+        self.mmdit = load_mmdit(
+            float16=True if self.dtype == self.float16_dtype else False,
+            key=self.mmdit_ckpt,
+            model_key=self.model_version,
+            low_memory_mode=self.low_memory_mode,
+            only_modulation_dict=only_modulation_dict,
+        )
+
+    def check_and_load_models(self):
+        if not hasattr(self, "mmdit"):
+            self.load_mmdit()
+        if not hasattr(self, "decoder"):
+            self.decoder = load_vae_decoder(
+                float16=True if self.dtype == self.float16_dtype else False,
+                key=self.mmdit_ckpt,
+            )
+        if not hasattr(self, "encoder"):
+            self.encoder = load_vae_encoder(float16=False, key=self.mmdit_ckpt)
+
+        if not hasattr(self, "clip_l"):
+            self.clip_l = load_text_encoder(
+                self.model,
+                float16=True if self.dtype == self.float16_dtype else False,
+                model_key="clip_l",
+            )
+            self.tokenizer_l = load_tokenizer(
+                self.model,
+                merges_key="tokenizer_l_merges",
+                vocab_key="tokenizer_l_vocab",
+                pad_with_eos=True,
+            )
+        if self.use_clip_g and not hasattr(self, "clip_g"):
+            self.clip_g = load_text_encoder(
+                self.model,
+                float16=True if self.dtype == self.float16_dtype else False,
+                model_key="clip_g",
+            )
+            self.tokenizer_g = load_tokenizer(
+                self.model,
+                merges_key="tokenizer_g_merges",
+                vocab_key="tokenizer_g_vocab",
+                pad_with_eos=False,
+            )
+        if self.use_t5 and not hasattr(self, "t5_encoder"):
             self.set_up_t5()
 
     def set_up_t5(self):
-        if self.t5_encoder is None:
+        if not hasattr(self, "t5_encoder") or self.t5_encoder is None:
             self.t5_encoder = load_t5_encoder(
                 float16=True if self.dtype == self.float16_dtype else False,
                 low_memory_mode=self.low_memory_mode,
             )
-        if self.t5_tokenizer is None:
+        if not hasattr(self, "t5_tokenizer") or self.t5_tokenizer is None:
             self.t5_tokenizer = load_t5_tokenizer()
         self.use_t5 = True
 
@@ -266,6 +290,7 @@ def generate_image(
         image_path: Optional[str] = None,
         denoise: float = 1.0,
     ):
+        self.check_and_load_models()
         # Start timing
         start_time = time.time()
 
@@ -405,10 +430,9 @@ def generate_image(
             )
             logger.info(f"Denoising time: {log['denoising']['time']}s")
 
-        # unload MMDIT and Sampler models after obtaining latents in low memory mode
+        # unload MMDIT model after obtaining latents in low memory mode
         if self.low_memory_mode:
             del self.mmdit
-            del self.sampler
             gc.collect()
 
         logger.debug(f"Latents dtype before casting: {latents.dtype}")
@@ -458,18 +482,19 @@ def generate_image(
                 f"Post decode active memory: {log['decoding']['post']['active_memory']}GB"
             )
 
-        logger.info("============= Summary =============")
-        logger.info(f"Text encoder: {log['text_encoding']['time']:.1f}s")
-        logger.info(f"Denoising: {log['denoising']['time']:.1f}s")
-        logger.info(f"Image decoder: {log['decoding']['time']:.1f}s")
-        logger.info(f"Peak memory: {log['peak_memory']:.1f}GB")
-
-        logger.info("============= Inference Context =============")
-        ic = DiffusionKitInferenceContext()
-        logger.info("Operating System:")
-        pprint(ic.os_spec())
-        logger.info("Device:")
-        pprint(ic.device_spec())
+        if verbose:
+            logger.info("============= Summary =============")
+            logger.info(f"Text encoder: {log['text_encoding']['time']:.1f}s")
+            logger.info(f"Denoising: {log['denoising']['time']:.1f}s")
+            logger.info(f"Image decoder: {log['decoding']['time']:.1f}s")
+            logger.info(f"Peak memory: {log['peak_memory']:.1f}GB")
+
+            logger.info("============= Inference Context =============")
+            ic = DiffusionKitInferenceContext()
+            logger.info("Operating System:")
+            pprint(ic.os_spec())
+            logger.info("Device:")
+            pprint(ic.device_spec())
 
         # unload VAE Decoder model after decoding in low memory mode
         if self.low_memory_mode:
@@ -566,30 +591,28 @@ def __init__(
         model_io._FLOAT16 = self.float16_dtype
         self.dtype = self.float16_dtype if w16 else mx.float32
         self.activation_dtype = self.float16_dtype if a16 else mx.float32
-        mmdit_ckpt = MMDIT_CKPT[model_version]
+        self.mmdit_ckpt = MMDIT_CKPT[model_version]
         self.low_memory_mode = low_memory_mode
-        self.mmdit = load_flux(float16=w16, low_memory_mode=low_memory_mode)
+        self.model = model
+        self.model_version = model_version
         self.sampler = FluxSampler(shift=shift)
-        self.decoder = load_vae_decoder(float16=w16, key=mmdit_ckpt)
-        self.encoder = load_vae_encoder(float16=False, key=mmdit_ckpt)
         self.latent_format = FluxLatentFormat()
         self.use_t5 = True
+        self.use_clip_g = False
+        self.check_and_load_models()
 
-        self.clip_l = load_text_encoder(
-            model,
-            w16,
-            model_key="clip_l",
-        )
-        self.tokenizer_l = load_tokenizer(
-            model,
-            merges_key="tokenizer_l_merges",
-            vocab_key="tokenizer_l_vocab",
-            pad_with_eos=True,
+    def load_mmdit(self, only_modulation_dict=False):
+        if only_modulation_dict:
+            return load_flux(
+                float16=True if self.dtype == self.float16_dtype else False,
+                low_memory_mode=self.low_memory_mode,
+                only_modulation_dict=only_modulation_dict,
+            )
+        self.mmdit = load_flux(
+            float16=True if self.dtype == self.float16_dtype else False,
+            low_memory_mode=self.low_memory_mode,
+            only_modulation_dict=only_modulation_dict,
         )
-        self.t5_encoder = None
-        self.t5_tokenizer = None
-        if self.use_t5:
-            self.set_up_t5()
 
     def encode_text(
         self,
@@ -634,7 +657,9 @@ def cache_modulation_params(self, pooled_text_embeddings, sigmas):
         )
 
     def clear_cache(self):
-        self.model.mmdit.clear_modulation_params_cache()
+        self.model.mmdit.load_weights(
+            self.model.load_mmdit(only_modulation_dict=True), strict=False
+        )
 
     def __call__(
         self,
@@ -731,6 +756,6 @@ def sample_euler(model: CFGDenoiser, x, sigmas, extra_args=None):
         end_time = t.format_dict["elapsed"]
         iter_time.append(round((end_time - start_time), 3))
 
-    # model.clear_cache()
+    model.clear_cache()
 
     return x, iter_time

python/src/diffusionkit/mlx/mmdit.py

@@ -10,6 +10,7 @@
 import numpy as np
 from argmaxtools.utils import get_logger
 from beartype.typing import Dict, List, Optional, Tuple
+from mlx.utils import tree_map
 
 from .config import MMDiTConfig, PositionalEncoding
 
@@ -163,7 +164,8 @@ def cache_modulation_params(
 
         self.to_offload = to_offload
         for x in self.to_offload:
-            x.clear()
+            x.update(tree_map(lambda _: mx.array([]), x.parameters()))
+            # x.clear()
 
         logger.info(f"Cached modulation_params for timesteps={timesteps}")
         logger.info(

python/src/diffusionkit/mlx/model_io.py

@@ -654,8 +654,10 @@ def load_mmdit(
     float16: bool = False,
     model_key: str = "mmdit_2b",
     low_memory_mode: bool = True,
+    only_modulation_dict: bool = False,
 ):
     """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.low_memory_mode = low_memory_mode
@@ -666,6 +668,9 @@ def load_mmdit(
     weights = mx.load(mmdit_weights_ckpt)
     weights = mmdit_state_dict_adjustments(weights, prefix="model.diffusion_model.")
     weights = {k: v.astype(dtype) for k, v in weights.items()}
+    if only_modulation_dict:
+        weights = {k: v for k, v in weights.items() if "adaLN" in k}
+        return tree_flatten(weights)
     model.update(tree_unflatten(tree_flatten(weights)))
 
     return model
@@ -676,6 +681,7 @@ def load_flux(
     float16: bool = False,
     model_key: str = "FLUX.1-schnell",
     low_memory_mode: bool = True,
+    only_modulation_dict: bool = False,
 ):
     """Load the MM-DiT Flux model from the checkpoint file."""
     dtype = _FLOAT16 if float16 else mx.float32
@@ -691,6 +697,9 @@ def load_flux(
         weights, prefix="", hidden_size=config.hidden_size, mlp_ratio=config.mlp_ratio
     )
     weights = {k: v.astype(dtype) for k, v in weights.items()}
+    if only_modulation_dict:
+        weights = {k: v for k, v in weights.items() if "adaLN" in k}
+        return tree_flatten(weights)
     model.update(tree_unflatten(tree_flatten(weights)))
 
     return model

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

@@ -112,10 +112,9 @@ def cli():
         args.cfg = 0.0
 
     if args.benchmark_mode:
-        raise NotImplementedError
-        # if args.low_memory_mode:
-        #     logger.warning("Benchmark mode is enabled, disabling low memory mode.")
-        # args.low_memory_mode = False
+        if args.low_memory_mode:
+            logger.warning("Benchmark mode is enabled, disabling low memory mode.")
+        args.low_memory_mode = False
 
     if args.denoise < 0.0 or args.denoise > 1.0:
         raise ValueError("Denoising factor must be between 0.0 and 1.0")