| 66 | self.vae = vae |
| 67 | |
| 68 | def simple_func(data_item): |
| 69 | image = self.transform(data_item["image"]) |
| 70 | latents = self.vae.encode(image.to(self.vae.device, dtype=self.vae.dtype).unsqueeze(0)).latent_dist.sample() |
| 71 | encoded_image = latents * self.vae.config.scaling_factor |
| 72 | encoded_image = encoded_image.detach() |
| 73 | encoded_image = encoded_image.squeeze(0).cpu() |
| 74 | |
| 75 | max_length = self.tokenizer.model_max_length |
| 76 | tokens = self.tokenizer( |
| 77 | [data_item["text"]], max_length=max_length, padding="max_length", truncation=True, return_tensors="pt" |
| 78 | ).input_ids |
| 79 | encoded_text = self.text_encoder(tokens.to(self.text_encoder.device))[0] |
| 80 | encoded_text = encoded_text.detach() |
| 81 | encoded_text = encoded_text.squeeze(0).cpu() |
| 82 | |
| 83 | return { |
| 84 | "image": encoded_image, |
| 85 | "text": encoded_text, |
| 86 | } |
| 87 | |
| 88 | self.pre_func = simple_func |
| 89 | |