| 280 | available_time_modes = ["all", "conv-only", "attn-only"] |
| 281 | |
| 282 | def __init__( |
| 283 | self, |
| 284 | *args, |
| 285 | video_kernel_size: Union[int, list] = 3, |
| 286 | alpha: float = 0.0, |
| 287 | merge_strategy: str = "learned", |
| 288 | time_mode: str = "conv-only", |
| 289 | **kwargs, |
| 290 | ): |
| 291 | self.video_kernel_size = video_kernel_size |
| 292 | self.alpha = alpha |
| 293 | self.merge_strategy = merge_strategy |
| 294 | self.time_mode = time_mode |
| 295 | assert ( |
| 296 | self.time_mode in self.available_time_modes |
| 297 | ), f"time_mode parameter has to be in {self.available_time_modes}" |
| 298 | super().__init__(*args, **kwargs) |
| 299 | |
| 300 | def get_last_layer(self, skip_time_mix=False, **kwargs): |
| 301 | if self.time_mode == "attn-only": |