(self, bs)
| 52 | raise ValueError(f"unknown merge strategy {self.merge_strategy}") |
| 53 | |
| 54 | def get_alpha(self, bs): |
| 55 | if self.merge_strategy == "fixed": |
| 56 | return self.mix_factor |
| 57 | elif self.merge_strategy == "learned": |
| 58 | return torch.sigmoid(self.mix_factor) |
| 59 | else: |
| 60 | raise NotImplementedError() |
| 61 | |
| 62 | def forward(self, x, temb, skip_video=False, timesteps=None): |
| 63 | if timesteps is None: |