| 349 | |
| 350 | |
| 351 | def encode_images(batch, image_vae, image_encoder, clip_image_processor, label_vae, image_column): |
| 352 | images = batch.pop("pixel_values") |
| 353 | pixel_values = torch.stack(list(images)) |
| 354 | pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float() |
| 355 | pixel_values = pixel_values.to(image_vae.device, dtype=image_vae.dtype) |
| 356 | |
| 357 | labels = batch.pop("label_pixel_values") |
| 358 | label_pixel_values = torch.stack(list(labels)) |
| 359 | label_pixel_values = label_pixel_values.to(memory_format=torch.contiguous_format).float() |
| 360 | label_pixel_values = label_pixel_values.to(label_vae.device, dtype=label_vae.dtype) |
| 361 | |
| 362 | clip_image = clip_image_processor.preprocess(batch[image_column], return_tensors="pt").data['pixel_values'] |
| 363 | clip_image = torch.stack(list(clip_image)) |
| 364 | clip_image = clip_image.to(memory_format=torch.contiguous_format).float() |
| 365 | clip_image = clip_image.to(image_encoder.device, dtype=image_encoder.dtype) |
| 366 | |
| 367 | with (torch.no_grad()): |
| 368 | image_latents = image_vae.encode(pixel_values).latent_dist.sample() |
| 369 | image_latents = image_latents * image_vae.config.scaling_factor |
| 370 | |
| 371 | label_latents = label_vae.encode(label_pixel_values).latent_dist.sample() |
| 372 | label_latents = label_latents * label_vae.config.scaling_factor |
| 373 | |
| 374 | image_embeds = image_encoder(clip_image).image_embeds |
| 375 | |
| 376 | return {"image_latents": image_latents.cpu(), "label_latents": label_latents.cpu(), "image_embeds": image_embeds.cpu()} |
| 377 | |
| 378 | |
| 379 | def make_train_dataset(args, accelerator): |