Spaces:
Paused
Paused
| import torch | |
| import os | |
| from src.datasets.text2image.caption_datasets import CaptionDataset | |
| from PIL import Image | |
| class BLIP3oDataset(CaptionDataset): | |
| def __getitem__(self, idx): | |
| if self.debug: | |
| idx = 0 | |
| try: | |
| data_sample = self.data_list[idx] | |
| if self.image_tokens_folder is not None: | |
| image_tokens = torch.load(os.path.join(self.image_tokens_folder, | |
| data_sample['image'] + '.pt')).long() | |
| data = dict(image_tokens=image_tokens) | |
| elif self.latents_ceph_folder is not None: | |
| image_latents = torch.load( | |
| self._read_ceph( | |
| os.path.join( | |
| self.latents_ceph_folder, data_sample['image'] + '.pt' | |
| ) | |
| ) | |
| ) | |
| data = dict(image_latents=image_latents) | |
| elif self.image_latents_folder is not None: | |
| image_latents = torch.load(os.path.join(self.image_latents_folder, | |
| data_sample['image'] + '.pt')) | |
| data = dict(image_latents=image_latents) | |
| else: | |
| if self.image_folder is not None: | |
| image = Image.open(os.path.join(self.image_folder,data_sample['image_path'])).convert('RGB') | |
| else: | |
| image = Image.open(data_sample['image_path']).convert('RGB') | |
| data = self._process_image(image) | |
| caption = data_sample['txt'] | |
| # print(caption) | |
| data["pixel_init"] = image | |
| data.update(self._process_text(caption)) | |
| data.update(image_dir=self.image_folder, image_file=None, | |
| type='text2image',text=caption) | |
| return data | |
| except Exception as e: | |
| print(f"Error when reading {self.data_path}:{self.data_list[idx]}: {e}", flush=True) | |
| return self._retry() | |