| 288 | |
| 289 | |
| 290 | class CLIPImageGridPointDiffusionTransformer(PointDiffusionTransformer): |
| 291 | def __init__( |
| 292 | self, |
| 293 | *, |
| 294 | device: torch.device, |
| 295 | dtype: torch.dtype, |
| 296 | n_ctx: int = 1024, |
| 297 | cond_drop_prob: float = 0.0, |
| 298 | frozen_clip: bool = True, |
| 299 | cache_dir: Optional[str] = None, |
| 300 | **kwargs, |
| 301 | ): |
| 302 | clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)( |
| 303 | device, |
| 304 | cache_dir=cache_dir, |
| 305 | ) |
| 306 | super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + clip.grid_size**2, **kwargs) |
| 307 | self.n_ctx = n_ctx |
| 308 | self.clip = clip |
| 309 | self.clip_embed = nn.Sequential( |
| 310 | nn.LayerNorm( |
| 311 | normalized_shape=(self.clip.grid_feature_dim,), device=device, dtype=dtype |
| 312 | ), |
| 313 | nn.Linear(self.clip.grid_feature_dim, self.backbone.width, device=device, dtype=dtype), |
| 314 | ) |
| 315 | self.cond_drop_prob = cond_drop_prob |
| 316 | |
| 317 | def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]) -> Dict[str, Any]: |
| 318 | _ = batch_size |
| 319 | with torch.no_grad(): |
| 320 | return dict(embeddings=self.clip.embed_images_grid(model_kwargs["images"])) |
| 321 | |
| 322 | def forward( |
| 323 | self, |
| 324 | x: torch.Tensor, |
| 325 | t: torch.Tensor, |
| 326 | images: Optional[Iterable[ImageType]] = None, |
| 327 | embeddings: Optional[Iterable[torch.Tensor]] = None, |
| 328 | ): |
| 329 | """ |
| 330 | :param x: an [N x C x T] tensor. |
| 331 | :param t: an [N] tensor. |
| 332 | :param images: a batch of images to condition on. |
| 333 | :param embeddings: a batch of CLIP latent grids to condition on. |
| 334 | :return: an [N x C' x T] tensor. |
| 335 | """ |
| 336 | assert images is not None or embeddings is not None, "must specify images or embeddings" |
| 337 | assert images is None or embeddings is None, "cannot specify both images and embeddings" |
| 338 | assert x.shape[-1] == self.n_ctx |
| 339 | |
| 340 | t_embed = self.time_embed(timestep_embedding(t, self.backbone.width)) |
| 341 | |
| 342 | if images is not None: |
| 343 | clip_out = self.clip.embed_images_grid(images) |
| 344 | else: |
| 345 | clip_out = embeddings |
| 346 | |
| 347 | if self.training: |