Apply suggestions from code review for flux dev

Author: ZachNagengast

python/src/diffusionkit/mlx/mmdit.py

@@ -214,15 +214,16 @@ def __call__(
         else:
             positional_encodings = None
 
-        timestep_embedding = self.guidance_in(self.t_embedder(timestep))
+        if config.guidance_embed:
+            timestep = self.guidance_in(self.t_embedder(timestep))
 
         # MultiModalTransformer layers
         if self.config.depth_multimodal > 0:
             for bidx, block in enumerate(self.multimodal_transformer_blocks):
                 latent_image_embeddings, token_level_text_embeddings = block(
                     latent_image_embeddings,
                     token_level_text_embeddings,
-                    timestep_embedding,
+                    timestep,
                     positional_encodings=positional_encodings,
                 )
 
@@ -245,7 +246,7 @@ def __call__(
 
         latent_image_embeddings = self.final_layer(
             latent_image_embeddings,
-            timestep_embedding
+            timestep,
         )
 
         if self.config.patchify_via_reshape:

python/src/diffusionkit/mlx/model_io.py

@@ -46,8 +46,8 @@
         "argmaxinc/mlx-FLUX.1-schnell-4bit-quantized": "flux-schnell-4bit-quantized.safetensors",
         "vae": "ae.safetensors",
     },
-    "raoulritter/flux-dev-mlx": {
-        "FLUX.1-dev": "flux1-dev-mlx.safetensors",
+    "argmaxinc/mlx-FLUX.1-dev": {
+        "argmaxinc/mlx-FLUX.1-dev": "flux1-dev.safetensors",
         "vae": "ae.safetensors",
     },
 }
@@ -79,7 +79,7 @@
         "vae_encoder": "encoder.",
         "vae_decoder": "decoder.",
     },
-    "raoulritter/flux-dev-mlx": {
+    "argmaxinc/mlx-FLUX.1-dev": {
             "vae_encoder": "encoder.",
             "vae_decoder": "decoder.",
         },