fixed nested if

argmaxinc · Aug 27, 2024 · d6d85bf · d6d85bf
1 parent 294ea0f
commit d6d85bf
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/python/src/diffusionkit/mlx/mmdit.py b/python/src/diffusionkit/mlx/mmdit.py
@@ -940,12 +940,11 @@ def affine_transform(
     norm_module: nn.Module = None,
 ) -> mx.array:
     """Affine transformation (Used for Adaptive LayerNorm Modulation)"""
-    if x.shape[0] == 1:
-        if norm_module is not None:
-            return mx.fast.layer_norm(
-                x, 1.0 + residual_scale.squeeze(), shift.squeeze(), norm_module.eps
-            )
-    if norm_module is not None:
+    if x.shape[0] == 1 and norm_module is not None:
+        return mx.fast.layer_norm(
+            x, 1.0 + residual_scale.squeeze(), shift.squeeze(), norm_module.eps
+        )
+    elif norm_module is not None:
         return norm_module(x) * (1.0 + residual_scale) + shift
     else:
         return x * (1.0 + residual_scale) + shift