(self, latents)
| 305 | |
| 306 | |
| 307 | def prepare_image_ids(self, latents): |
| 308 | batch_size, _, height, width = latents.shape |
| 309 | latent_image_ids = torch.zeros(height // 2, width // 2, 3) |
| 310 | latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height // 2)[:, None] |
| 311 | latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width // 2)[None, :] |
| 312 | |
| 313 | latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape |
| 314 | |
| 315 | latent_image_ids = latent_image_ids[None, :].repeat(batch_size, 1, 1, 1) |
| 316 | latent_image_ids = latent_image_ids.reshape( |
| 317 | batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels |
| 318 | ) |
| 319 | latent_image_ids = latent_image_ids.to(device=latents.device, dtype=latents.dtype) |
| 320 | |
| 321 | return latent_image_ids |
| 322 | |
| 323 | |
| 324 | def tiled_forward( |
no test coverage detected