iyedjb/self-forcing

1from utils.lmdb import get_array_shape_from_lmdb, retrieve_row_from_lmdb2from torch.utils.data import Dataset3import numpy as np4import torch5import lmdb6import json7from pathlib import Path8from PIL import Image9import os10 11 12class TextDataset(Dataset):13    def __init__(self, prompt_path, extended_prompt_path=None):14        with open(prompt_path, encoding="utf-8") as f:15            self.prompt_list = [line.rstrip() for line in f]16 17        if extended_prompt_path is not None:18            with open(extended_prompt_path, encoding="utf-8") as f:19                self.extended_prompt_list = [line.rstrip() for line in f]20            assert len(self.extended_prompt_list) == len(self.prompt_list)21        else:22            self.extended_prompt_list = None23 24    def __len__(self):25        return len(self.prompt_list)26 27    def __getitem__(self, idx):28        batch = {29            "prompts": self.prompt_list[idx],30            "idx": idx,31        }32        if self.extended_prompt_list is not None:33            batch["extended_prompts"] = self.extended_prompt_list[idx]34        return batch35 36 37class ODERegressionLMDBDataset(Dataset):38    def __init__(self, data_path: str, max_pair: int = int(1e8)):39        self.env = lmdb.open(data_path, readonly=True,40                             lock=False, readahead=False, meminit=False)41 42        self.latents_shape = get_array_shape_from_lmdb(self.env, 'latents')43        self.max_pair = max_pair44 45    def __len__(self):46        return min(self.latents_shape[0], self.max_pair)47 48    def __getitem__(self, idx):49        """50        Outputs:51            - prompts: List of Strings52            - latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.53        """54        latents = retrieve_row_from_lmdb(55            self.env,56            "latents", np.float16, idx, shape=self.latents_shape[1:]57        )58 59        if len(latents.shape) == 4:60            latents = latents[None, ...]61 62        prompts = retrieve_row_from_lmdb(63            self.env,64            "prompts", str, idx65        )66        return {67            "prompts": prompts,68            "ode_latent": torch.tensor(latents, dtype=torch.float32)69        }70 71 72class ShardingLMDBDataset(Dataset):73    def __init__(self, data_path: str, max_pair: int = int(1e8)):74        self.envs = []75        self.index = []76 77        for fname in sorted(os.listdir(data_path)):78            path = os.path.join(data_path, fname)79            env = lmdb.open(path,80                            readonly=True,81                            lock=False,82                            readahead=False,83                            meminit=False)84            self.envs.append(env)85 86        self.latents_shape = [None] * len(self.envs)87        for shard_id, env in enumerate(self.envs):88            self.latents_shape[shard_id] = get_array_shape_from_lmdb(env, 'latents')89            for local_i in range(self.latents_shape[shard_id][0]):90                self.index.append((shard_id, local_i))91 92            # print("shard_id ", shard_id, " local_i ", local_i)93 94        self.max_pair = max_pair95 96    def __len__(self):97        return len(self.index)98 99    def __getitem__(self, idx):100        """101            Outputs:102                - prompts: List of Strings103                - latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.104        """105        shard_id, local_idx = self.index[idx]106 107        latents = retrieve_row_from_lmdb(108            self.envs[shard_id],109            "latents", np.float16, local_idx,110            shape=self.latents_shape[shard_id][1:]111        )112 113        if len(latents.shape) == 4:114            latents = latents[None, ...]115 116        prompts = retrieve_row_from_lmdb(117            self.envs[shard_id],118            "prompts", str, local_idx119        )120 121        return {122            "prompts": prompts,123            "ode_latent": torch.tensor(latents, dtype=torch.float32)124        }125 126 127class TextImagePairDataset(Dataset):128    def __init__(129        self,130        data_dir,131        transform=None,132        eval_first_n=-1,133        pad_to_multiple_of=None134    ):135        """136        Args:137            data_dir (str): Path to the directory containing:138                - target_crop_info_*.json (metadata file)139                - */ (subdirectory containing images with matching aspect ratio)140            transform (callable, optional): Optional transform to be applied on the image141        """142        self.transform = transform143        data_dir = Path(data_dir)144 145        # Find the metadata JSON file146        metadata_files = list(data_dir.glob('target_crop_info_*.json'))147        if not metadata_files:148            raise FileNotFoundError(f"No metadata file found in {data_dir}")149        if len(metadata_files) > 1:150            raise ValueError(f"Multiple metadata files found in {data_dir}")151 152        metadata_path = metadata_files[0]153        # Extract aspect ratio from metadata filename (e.g. target_crop_info_26-15.json -> 26-15)154        aspect_ratio = metadata_path.stem.split('_')[-1]155 156        # Use aspect ratio subfolder for images157        self.image_dir = data_dir / aspect_ratio158        if not self.image_dir.exists():159            raise FileNotFoundError(f"Image directory not found: {self.image_dir}")160 161        # Load metadata162        with open(metadata_path, 'r') as f:163            self.metadata = json.load(f)164 165        eval_first_n = eval_first_n if eval_first_n != -1 else len(self.metadata)166        self.metadata = self.metadata[:eval_first_n]167 168        # Verify all images exist169        for item in self.metadata:170            image_path = self.image_dir / item['file_name']171            if not image_path.exists():172                raise FileNotFoundError(f"Image not found: {image_path}")173 174        self.dummy_prompt = "DUMMY PROMPT"175        self.pre_pad_len = len(self.metadata)176        if pad_to_multiple_of is not None and len(self.metadata) % pad_to_multiple_of != 0:177            # Duplicate the last entry178            self.metadata += [self.metadata[-1]] * (179                pad_to_multiple_of - len(self.metadata) % pad_to_multiple_of180            )181 182    def __len__(self):183        return len(self.metadata)184 185    def __getitem__(self, idx):186        """187        Returns:188            dict: A dictionary containing:189                - image: PIL Image190                - caption: str191                - target_bbox: list of int [x1, y1, x2, y2]192                - target_ratio: str193                - type: str194                - origin_size: tuple of int (width, height)195        """196        item = self.metadata[idx]197 198        # Load image199        image_path = self.image_dir / item['file_name']200        image = Image.open(image_path).convert('RGB')201 202        # Apply transform if specified203        if self.transform:204            image = self.transform(image)205 206        return {207            'image': image,208            'prompts': item['caption'],209            'target_bbox': item['target_crop']['target_bbox'],210            'target_ratio': item['target_crop']['target_ratio'],211            'type': item['type'],212            'origin_size': (item['origin_width'], item['origin_height']),213            'idx': idx214        }215 216 217def cycle(dl):218    while True:219        for data in dl:220            yield data221