| 123 | self.max_pwatermark = max_pwatermark # filter out watermarked images |
| 124 | |
| 125 | def make_loader(self, dataset_config, train=True): |
| 126 | if 'image_transforms' in dataset_config: |
| 127 | image_transforms = [instantiate_from_config(tt) for tt in dataset_config.image_transforms] |
| 128 | else: |
| 129 | image_transforms = [] |
| 130 | |
| 131 | image_transforms.extend([torchvision.transforms.ToTensor(), |
| 132 | torchvision.transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))]) |
| 133 | image_transforms = torchvision.transforms.Compose(image_transforms) |
| 134 | |
| 135 | if 'transforms' in dataset_config: |
| 136 | transforms_config = OmegaConf.to_container(dataset_config.transforms) |
| 137 | else: |
| 138 | transforms_config = dict() |
| 139 | |
| 140 | transform_dict = {dkey: load_partial_from_config(transforms_config[dkey]) |
| 141 | if transforms_config[dkey] != 'identity' else identity |
| 142 | for dkey in transforms_config} |
| 143 | img_key = dataset_config.get('image_key', 'jpeg') |
| 144 | transform_dict.update({img_key: image_transforms}) |
| 145 | |
| 146 | if 'postprocess' in dataset_config: |
| 147 | postprocess = instantiate_from_config(dataset_config['postprocess']) |
| 148 | else: |
| 149 | postprocess = None |
| 150 | |
| 151 | shuffle = dataset_config.get('shuffle', 0) |
| 152 | shardshuffle = shuffle > 0 |
| 153 | |
| 154 | nodesplitter = wds.shardlists.split_by_node if self.multinode else wds.shardlists.single_node_only |
| 155 | |
| 156 | if self.tar_base == "__improvedaesthetic__": |
| 157 | print("## Warning, loading the same improved aesthetic dataset " |
| 158 | "for all splits and ignoring shards parameter.") |
| 159 | tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -" |
| 160 | else: |
| 161 | tars = os.path.join(self.tar_base, dataset_config.shards) |
| 162 | |
| 163 | dset = wds.WebDataset( |
| 164 | tars, |
| 165 | nodesplitter=nodesplitter, |
| 166 | shardshuffle=shardshuffle, |
| 167 | handler=wds.warn_and_continue).repeat().shuffle(shuffle) |
| 168 | print(f'Loading webdataset with {len(dset.pipeline[0].urls)} shards.') |
| 169 | |
| 170 | dset = (dset |
| 171 | .select(self.filter_keys) |
| 172 | .decode('pil', handler=wds.warn_and_continue) |
| 173 | .select(self.filter_size) |
| 174 | .map_dict(**transform_dict, handler=wds.warn_and_continue) |
| 175 | ) |
| 176 | if postprocess is not None: |
| 177 | dset = dset.map(postprocess) |
| 178 | dset = (dset |
| 179 | .batched(self.batch_size, partial=False, |
| 180 | collation_fn=dict_collation_fn) |
| 181 | ) |
| 182 | |