(self, x, timesteps, skip_video=False)
| 134 | raise ValueError(f"unknown merge strategy {self.merge_strategy}") |
| 135 | |
| 136 | def forward(self, x, timesteps, skip_video=False): |
| 137 | if skip_video: |
| 138 | return super().forward(x) |
| 139 | |
| 140 | x_in = x |
| 141 | x = self.attention(x) |
| 142 | h, w = x.shape[2:] |
| 143 | x = rearrange(x, "b c h w -> b (h w) c") |
| 144 | |
| 145 | x_mix = x |
| 146 | num_frames = torch.arange(timesteps, device=x.device) |
| 147 | num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps) |
| 148 | num_frames = rearrange(num_frames, "b t -> (b t)") |
| 149 | t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False) |
| 150 | emb = self.video_time_embed(t_emb) # b, n_channels |
| 151 | emb = emb[:, None, :] |
| 152 | x_mix = x_mix + emb |
| 153 | |
| 154 | alpha = self.get_alpha() |
| 155 | x_mix = self.time_mix_block(x_mix, timesteps=timesteps) |
| 156 | x = alpha * x + (1.0 - alpha) * x_mix # alpha merge |
| 157 | |
| 158 | x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w) |
| 159 | x = self.proj_out(x) |
| 160 | |
| 161 | return x_in + x |
| 162 | |
| 163 | def get_alpha( |
| 164 | self, |
nothing calls this directly
no test coverage detected