Merge pull request #25 from argmaxinc/sd3-fix

Fix sd3 layer norm bug
argmaxinc · Aug 27, 2024 · d915340 · d915340
2 parents ccf6509 + ad0ad24
commit d915340
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/python/src/diffusionkit/mlx/mmdit.py b/python/src/diffusionkit/mlx/mmdit.py
@@ -940,11 +940,14 @@ def affine_transform(
     norm_module: nn.Module = None,
 ) -> mx.array:
     """Affine transformation (Used for Adaptive LayerNorm Modulation)"""
-    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
         )
-    return x * (1.0 + residual_scale) + shift
+    elif norm_module is not None:
+        return norm_module(x) * (1.0 + residual_scale) + shift
+    else:
+        return x * (1.0 + residual_scale) + shift
 
 
 def unpatchify(

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 from setuptools import find_packages, setup
 from setuptools.command.install import install
 
-VERSION = "0.3.2"
+VERSION = "0.3.3"
 
 
 class VersionInstallCommand(install):