| 21 | |
| 22 | |
| 23 | class CogAdaLayerNorm(torch.nn.Module): |
| 24 | def __init__(self, dim, dim_cond, single=False): |
| 25 | super().__init__() |
| 26 | self.single = single |
| 27 | self.linear = torch.nn.Linear(dim_cond, dim * (2 if single else 6)) |
| 28 | self.norm = torch.nn.LayerNorm(dim, elementwise_affine=True, eps=1e-5) |
| 29 | |
| 30 | |
| 31 | def forward(self, hidden_states, prompt_emb, emb): |
| 32 | emb = self.linear(torch.nn.functional.silu(emb)) |
| 33 | if self.single: |
| 34 | shift, scale = emb.unsqueeze(1).chunk(2, dim=2) |
| 35 | hidden_states = self.norm(hidden_states) * (1 + scale) + shift |
| 36 | return hidden_states |
| 37 | else: |
| 38 | shift_a, scale_a, gate_a, shift_b, scale_b, gate_b = emb.unsqueeze(1).chunk(6, dim=2) |
| 39 | hidden_states = self.norm(hidden_states) * (1 + scale_a) + shift_a |
| 40 | prompt_emb = self.norm(prompt_emb) * (1 + scale_b) + shift_b |
| 41 | return hidden_states, prompt_emb, gate_a, gate_b |
| 42 | |
| 43 | |
| 44 | |