| 356 | |
| 357 | |
| 358 | class UpsamplePointDiffusionTransformer(PointDiffusionTransformer): |
| 359 | def __init__( |
| 360 | self, |
| 361 | *, |
| 362 | device: torch.device, |
| 363 | dtype: torch.dtype, |
| 364 | cond_input_channels: Optional[int] = None, |
| 365 | cond_ctx: int = 1024, |
| 366 | n_ctx: int = 4096 - 1024, |
| 367 | channel_scales: Optional[Sequence[float]] = None, |
| 368 | channel_biases: Optional[Sequence[float]] = None, |
| 369 | **kwargs, |
| 370 | ): |
| 371 | super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + cond_ctx, **kwargs) |
| 372 | self.n_ctx = n_ctx |
| 373 | self.cond_input_channels = cond_input_channels or self.input_channels |
| 374 | self.cond_point_proj = nn.Linear( |
| 375 | self.cond_input_channels, self.backbone.width, device=device, dtype=dtype |
| 376 | ) |
| 377 | |
| 378 | self.register_buffer( |
| 379 | "channel_scales", |
| 380 | torch.tensor(channel_scales, dtype=dtype, device=device) |
| 381 | if channel_scales is not None |
| 382 | else None, |
| 383 | ) |
| 384 | self.register_buffer( |
| 385 | "channel_biases", |
| 386 | torch.tensor(channel_biases, dtype=dtype, device=device) |
| 387 | if channel_biases is not None |
| 388 | else None, |
| 389 | ) |
| 390 | |
| 391 | def forward(self, x: torch.Tensor, t: torch.Tensor, *, low_res: torch.Tensor): |
| 392 | """ |
| 393 | :param x: an [N x C1 x T] tensor. |
| 394 | :param t: an [N] tensor. |
| 395 | :param low_res: an [N x C2 x T'] tensor of conditioning points. |
| 396 | :return: an [N x C3 x T] tensor. |
| 397 | """ |
| 398 | assert x.shape[-1] == self.n_ctx |
| 399 | t_embed = self.time_embed(timestep_embedding(t, self.backbone.width)) |
| 400 | low_res_embed = self._embed_low_res(low_res) |
| 401 | cond = [(t_embed, self.time_token_cond), (low_res_embed, True)] |
| 402 | return self._forward_with_cond(x, cond) |
| 403 | |
| 404 | def _embed_low_res(self, x: torch.Tensor) -> torch.Tensor: |
| 405 | if self.channel_scales is not None: |
| 406 | x = x * self.channel_scales[None, :, None] |
| 407 | if self.channel_biases is not None: |
| 408 | x = x + self.channel_biases[None, :, None] |
| 409 | return self.cond_point_proj(x.permute(0, 2, 1)) |
| 410 | |
| 411 | |
| 412 | class CLIPImageGridUpsamplePointDiffusionTransformer(UpsamplePointDiffusionTransformer): |