(self, input, timesteps, skip_video=False)
| 97 | ) |
| 98 | |
| 99 | def forward(self, input, timesteps, skip_video=False): |
| 100 | x = super().forward(input) |
| 101 | if skip_video: |
| 102 | return x |
| 103 | x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps) |
| 104 | x = self.time_mix_conv(x) |
| 105 | return rearrange(x, "b c t h w -> (b t) c h w") |
| 106 | |
| 107 | |
| 108 | class VideoBlock(AttnBlock): |