Update app.py

app.py

@@ -144,15 +144,16 @@ class DDiTBlock(nn.Module):
         self.attn = SelfAttention(config)
         self.ln_2 = nn.LayerNorm(config.n_embd, bias=config.bias)
         self.mlp = MLP(config)
-        self.adaLN_modulation = nn.Linear(config.cond_dim, 6 * config.n_embd, bias=True)
+
+        self.adaLN_modulation = nn.Linear(config.cond_dim, 6 * config.n_embd)
         self.adaLN_modulation.weight.data.zero_()
         self.adaLN_modulation.bias.data.zero_()
+
     def forward(self, x, c):
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
-        x_skip = x
-        x = modulate(self.ln_1(x), shift_msa, scale_msa)
-        x = self.attn(x)
-        x = bias_add_scale(x, None, gate_msa, x_skip)
+        x_skip = x        
+        modulated_x = modulate(self.ln_1(x), shift_msa, scale_msa)
+        x = bias_add_scale(self.attn(self.ln_1(x)), None, gate_msa, x_skip)
         x = bias_add_scale(self.mlp(modulate(self.ln_2(x), shift_mlp, scale_mlp)), None, gate_mlp, x)
         return x