| 60 | raise NotImplementedError() |
| 61 | |
| 62 | def forward(self, x, temb, skip_video=False, timesteps=None): |
| 63 | if timesteps is None: |
| 64 | timesteps = self.timesteps |
| 65 | |
| 66 | b, c, h, w = x.shape |
| 67 | |
| 68 | x = super().forward(x, temb) |
| 69 | |
| 70 | if not skip_video: |
| 71 | x_mix = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps) |
| 72 | |
| 73 | x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps) |
| 74 | |
| 75 | x = self.time_stack(x, temb) |
| 76 | |
| 77 | alpha = self.get_alpha(bs=b // timesteps) |
| 78 | x = alpha * x + (1.0 - alpha) * x_mix |
| 79 | |
| 80 | x = rearrange(x, "b c t h w -> (b t) c h w") |
| 81 | return x |
| 82 | |
| 83 | |
| 84 | class AE3DConv(torch.nn.Conv2d): |