(
self,
x: torch.Tensor,
context: Optional[torch.Tensor] = None,
time_context: Optional[torch.Tensor] = None,
timesteps: Optional[int] = None,
image_only_indicator: Optional[torch.Tensor] = None,
)
| 225 | self.dtype = str_to_dtype[dtype] |
| 226 | |
| 227 | def forward( |
| 228 | self, |
| 229 | x: torch.Tensor, |
| 230 | context: Optional[torch.Tensor] = None, |
| 231 | time_context: Optional[torch.Tensor] = None, |
| 232 | timesteps: Optional[int] = None, |
| 233 | image_only_indicator: Optional[torch.Tensor] = None, |
| 234 | ) -> torch.Tensor: |
| 235 | _, _, h, w = x.shape |
| 236 | x_in = x |
| 237 | spatial_context = None |
| 238 | if exists(context): |
| 239 | spatial_context = context |
| 240 | |
| 241 | if self.use_spatial_context: |
| 242 | assert context.ndim == 3, f"n dims of spatial context should be 3 but are {context.ndim}" |
| 243 | |
| 244 | time_context = context |
| 245 | time_context_first_timestep = time_context[::timesteps] |
| 246 | time_context = repeat(time_context_first_timestep, "b ... -> (b n) ...", n=h * w) |
| 247 | elif time_context is not None and not self.use_spatial_context: |
| 248 | time_context = repeat(time_context, "b ... -> (b n) ...", n=h * w) |
| 249 | if time_context.ndim == 2: |
| 250 | time_context = rearrange(time_context, "b c -> b 1 c") |
| 251 | |
| 252 | x = self.norm(x) |
| 253 | if not self.use_linear: |
| 254 | x = self.proj_in(x) |
| 255 | x = rearrange(x, "b c h w -> b (h w) c") |
| 256 | if self.use_linear: |
| 257 | x = self.proj_in(x) |
| 258 | |
| 259 | num_frames = torch.arange(timesteps, device=x.device) |
| 260 | num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps) |
| 261 | num_frames = rearrange(num_frames, "b t -> (b t)") |
| 262 | t_emb = timestep_embedding( |
| 263 | num_frames, |
| 264 | self.in_channels, |
| 265 | repeat_only=False, |
| 266 | max_period=self.max_time_embed_period, |
| 267 | dtype=self.dtype, |
| 268 | ) |
| 269 | emb = self.time_pos_embed(t_emb) |
| 270 | emb = emb[:, None, :] |
| 271 | |
| 272 | for it_, (block, mix_block) in enumerate(zip(self.transformer_blocks, self.time_stack)): |
| 273 | x = block( |
| 274 | x, |
| 275 | context=spatial_context, |
| 276 | ) |
| 277 | |
| 278 | x_mix = x |
| 279 | x_mix = x_mix + emb |
| 280 | |
| 281 | x_mix = mix_block(x_mix, context=time_context, timesteps=timesteps) |
| 282 | x = self.time_mixer( |
| 283 | x_spatial=x, |
| 284 | x_temporal=x_mix, |
nothing calls this directly
no test coverage detected