Merge pull request #7 from LightricksResearch/feature/fix-transformer-init-bug

xora/models/transformers/transformer3d.py

@@ -186,14 +186,14 @@ class Transformer3DModel(ModelMixin, ConfigMixin):
 
         # Zero-out adaLN modulation layers in PixArt blocks:
         for block in self.transformer_blocks:
-            if mode == "xora":
+            if mode.lower() == "xora":
                 nn.init.constant_(block.attn1.to_out[0].weight, 0)
                 nn.init.constant_(block.attn1.to_out[0].bias, 0)
 
             nn.init.constant_(block.attn2.to_out[0].weight, 0)
             nn.init.constant_(block.attn2.to_out[0].bias, 0)
 
-            if mode == "xora":
+            if mode.lower() == "xora":
                 nn.init.constant_(block.ff.net[2].weight, 0)
                 nn.init.constant_(block.ff.net[2].bias, 0)