| 103 | return self._forward(x, context, timesteps=timesteps) |
| 104 | |
| 105 | def _forward(self, x, context=None, timesteps=None): |
| 106 | assert self.timesteps or timesteps |
| 107 | assert not (self.timesteps and timesteps) or self.timesteps == timesteps |
| 108 | timesteps = self.timesteps or timesteps |
| 109 | B, S, C = x.shape |
| 110 | x = rearrange(x, "(b t) s c -> (b s) t c", t=timesteps) |
| 111 | |
| 112 | if self.ff_in: |
| 113 | x_skip = x |
| 114 | x = self.ff_in(self.norm_in(x)) |
| 115 | if self.is_res: |
| 116 | x += x_skip |
| 117 | |
| 118 | if self.disable_self_attn: |
| 119 | x = self.attn1(self.norm1(x), context=context) + x |
| 120 | else: |
| 121 | x = self.attn1(self.norm1(x)) + x |
| 122 | |
| 123 | if self.attn2 is not None: |
| 124 | if self.switch_temporal_ca_to_sa: |
| 125 | x = self.attn2(self.norm2(x)) + x |
| 126 | else: |
| 127 | x = self.attn2(self.norm2(x), context=context) + x |
| 128 | x_skip = x |
| 129 | x = self.ff(self.norm3(x)) |
| 130 | if self.is_res: |
| 131 | x += x_skip |
| 132 | |
| 133 | x = rearrange(x, "(b s) t c -> (b t) s c", s=S, b=B // timesteps, c=C, t=timesteps) |
| 134 | return x |
| 135 | |
| 136 | def get_last_layer(self): |
| 137 | return self.ff.net[-1].weight |