Spaces:
Sleeping
Sleeping
import webdataset as wds | |
import kornia | |
from PIL import Image | |
import io | |
import os | |
import torchvision | |
from PIL import Image | |
import glob | |
import random | |
import numpy as np | |
import pytorch_lightning as pl | |
from tqdm import tqdm | |
from omegaconf import OmegaConf | |
from einops import rearrange | |
import torch | |
from webdataset.handlers import warn_and_continue | |
from ldm.util import instantiate_from_config | |
from ldm.data.inpainting.synthetic_mask import gen_large_mask, MASK_MODES | |
from ldm.data.base import PRNGMixin | |
class DataWithWings(torch.utils.data.IterableDataset): | |
def __init__(self, min_size, transform=None, target_transform=None): | |
self.min_size = min_size | |
self.transform = transform if transform is not None else nn.Identity() | |
self.target_transform = target_transform if target_transform is not None else nn.Identity() | |
self.kv = OnDiskKV(file='/home/ubuntu/laion5B-watermark-safety-ordered', key_format='q', value_format='ee') | |
self.kv_aesthetic = OnDiskKV(file='/home/ubuntu/laion5B-aesthetic-tags-kv', key_format='q', value_format='e') | |
self.pwatermark_threshold = 0.8 | |
self.punsafe_threshold = 0.5 | |
self.aesthetic_threshold = 5. | |
self.total_samples = 0 | |
self.samples = 0 | |
location = 'pipe:aws s3 cp --quiet s3://s-datasets/laion5b/laion2B-data/{000000..231349}.tar -' | |
self.inner_dataset = wds.DataPipeline( | |
wds.ResampledShards(location), | |
wds.tarfile_to_samples(handler=wds.warn_and_continue), | |
wds.shuffle(1000, handler=wds.warn_and_continue), | |
wds.decode('pilrgb', handler=wds.warn_and_continue), | |
wds.map(self._add_tags, handler=wds.ignore_and_continue), | |
wds.select(self._filter_predicate), | |
wds.map_dict(jpg=self.transform, txt=self.target_transform, punsafe=self._punsafe_to_class, handler=wds.warn_and_continue), | |
wds.to_tuple('jpg', 'txt', 'punsafe', handler=wds.warn_and_continue), | |
) | |
def _compute_hash(url, text): | |
if url is None: | |
url = '' | |
if text is None: | |
text = '' | |
total = (url + text).encode('utf-8') | |
return mmh3.hash64(total)[0] | |
def _add_tags(self, x): | |
hsh = self._compute_hash(x['json']['url'], x['txt']) | |
pwatermark, punsafe = self.kv[hsh] | |
aesthetic = self.kv_aesthetic[hsh][0] | |
return {**x, 'pwatermark': pwatermark, 'punsafe': punsafe, 'aesthetic': aesthetic} | |
def _punsafe_to_class(self, punsafe): | |
return torch.tensor(punsafe >= self.punsafe_threshold).long() | |
def _filter_predicate(self, x): | |
try: | |
return x['pwatermark'] < self.pwatermark_threshold and x['aesthetic'] >= self.aesthetic_threshold and x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size | |
except: | |
return False | |
def __iter__(self): | |
return iter(self.inner_dataset) | |
def dict_collation_fn(samples, combine_tensors=True, combine_scalars=True): | |
"""Take a list of samples (as dictionary) and create a batch, preserving the keys. | |
If `tensors` is True, `ndarray` objects are combined into | |
tensor batches. | |
:param dict samples: list of samples | |
:param bool tensors: whether to turn lists of ndarrays into a single ndarray | |
:returns: single sample consisting of a batch | |
:rtype: dict | |
""" | |
keys = set.intersection(*[set(sample.keys()) for sample in samples]) | |
batched = {key: [] for key in keys} | |
for s in samples: | |
[batched[key].append(s[key]) for key in batched] | |
result = {} | |
for key in batched: | |
if isinstance(batched[key][0], (int, float)): | |
if combine_scalars: | |
result[key] = np.array(list(batched[key])) | |
elif isinstance(batched[key][0], torch.Tensor): | |
if combine_tensors: | |
result[key] = torch.stack(list(batched[key])) | |
elif isinstance(batched[key][0], np.ndarray): | |
if combine_tensors: | |
result[key] = np.array(list(batched[key])) | |
else: | |
result[key] = list(batched[key]) | |
return result | |
class WebDataModuleFromConfig(pl.LightningDataModule): | |
def __init__(self, tar_base, batch_size, train=None, validation=None, | |
test=None, num_workers=4, multinode=True, min_size=None, | |
max_pwatermark=1.0, | |
**kwargs): | |
super().__init__(self) | |
print(f'Setting tar base to {tar_base}') | |
self.tar_base = tar_base | |
self.batch_size = batch_size | |
self.num_workers = num_workers | |
self.train = train | |
self.validation = validation | |
self.test = test | |
self.multinode = multinode | |
self.min_size = min_size # filter out very small images | |
self.max_pwatermark = max_pwatermark # filter out watermarked images | |
def make_loader(self, dataset_config, train=True): | |
if 'image_transforms' in dataset_config: | |
image_transforms = [instantiate_from_config(tt) for tt in dataset_config.image_transforms] | |
else: | |
image_transforms = [] | |
image_transforms.extend([torchvision.transforms.ToTensor(), | |
torchvision.transforms.Lambda(lambda x: rearrange(x * 2. - 1., 'c h w -> h w c'))]) | |
image_transforms = torchvision.transforms.Compose(image_transforms) | |
if 'transforms' in dataset_config: | |
transforms_config = OmegaConf.to_container(dataset_config.transforms) | |
else: | |
transforms_config = dict() | |
transform_dict = {dkey: load_partial_from_config(transforms_config[dkey]) | |
if transforms_config[dkey] != 'identity' else identity | |
for dkey in transforms_config} | |
img_key = dataset_config.get('image_key', 'jpeg') | |
transform_dict.update({img_key: image_transforms}) | |
if 'postprocess' in dataset_config: | |
postprocess = instantiate_from_config(dataset_config['postprocess']) | |
else: | |
postprocess = None | |
shuffle = dataset_config.get('shuffle', 0) | |
shardshuffle = shuffle > 0 | |
nodesplitter = wds.shardlists.split_by_node if self.multinode else wds.shardlists.single_node_only | |
if self.tar_base == "__improvedaesthetic__": | |
print("## Warning, loading the same improved aesthetic dataset " | |
"for all splits and ignoring shards parameter.") | |
tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -" | |
else: | |
tars = os.path.join(self.tar_base, dataset_config.shards) | |
dset = wds.WebDataset( | |
tars, | |
nodesplitter=nodesplitter, | |
shardshuffle=shardshuffle, | |
handler=wds.warn_and_continue).repeat().shuffle(shuffle) | |
print(f'Loading webdataset with {len(dset.pipeline[0].urls)} shards.') | |
dset = (dset | |
.select(self.filter_keys) | |
.decode('pil', handler=wds.warn_and_continue) | |
.select(self.filter_size) | |
.map_dict(**transform_dict, handler=wds.warn_and_continue) | |
) | |
if postprocess is not None: | |
dset = dset.map(postprocess) | |
dset = (dset | |
.batched(self.batch_size, partial=False, | |
collation_fn=dict_collation_fn) | |
) | |
loader = wds.WebLoader(dset, batch_size=None, shuffle=False, | |
num_workers=self.num_workers) | |
return loader | |
def filter_size(self, x): | |
try: | |
valid = True | |
if self.min_size is not None and self.min_size > 1: | |
try: | |
valid = valid and x['json']['original_width'] >= self.min_size and x['json']['original_height'] >= self.min_size | |
except Exception: | |
valid = False | |
if self.max_pwatermark is not None and self.max_pwatermark < 1.0: | |
try: | |
valid = valid and x['json']['pwatermark'] <= self.max_pwatermark | |
except Exception: | |
valid = False | |
return valid | |
except Exception: | |
return False | |
def filter_keys(self, x): | |
try: | |
return ("jpg" in x) and ("txt" in x) | |
except Exception: | |
return False | |
def train_dataloader(self): | |
return self.make_loader(self.train) | |
def val_dataloader(self): | |
return self.make_loader(self.validation, train=False) | |
def test_dataloader(self): | |
return self.make_loader(self.test, train=False) | |
from ldm.modules.image_degradation import degradation_fn_bsr_light | |
import cv2 | |
class AddLR(object): | |
def __init__(self, factor, output_size, initial_size=None, image_key="jpg"): | |
self.factor = factor | |
self.output_size = output_size | |
self.image_key = image_key | |
self.initial_size = initial_size | |
def pt2np(self, x): | |
x = ((x+1.0)*127.5).clamp(0, 255).to(dtype=torch.uint8).detach().cpu().numpy() | |
return x | |
def np2pt(self, x): | |
x = torch.from_numpy(x)/127.5-1.0 | |
return x | |
def __call__(self, sample): | |
# sample['jpg'] is tensor hwc in [-1, 1] at this point | |
x = self.pt2np(sample[self.image_key]) | |
if self.initial_size is not None: | |
x = cv2.resize(x, (self.initial_size, self.initial_size), interpolation=2) | |
x = degradation_fn_bsr_light(x, sf=self.factor)['image'] | |
x = cv2.resize(x, (self.output_size, self.output_size), interpolation=2) | |
x = self.np2pt(x) | |
sample['lr'] = x | |
return sample | |
class AddBW(object): | |
def __init__(self, image_key="jpg"): | |
self.image_key = image_key | |
def pt2np(self, x): | |
x = ((x+1.0)*127.5).clamp(0, 255).to(dtype=torch.uint8).detach().cpu().numpy() | |
return x | |
def np2pt(self, x): | |
x = torch.from_numpy(x)/127.5-1.0 | |
return x | |
def __call__(self, sample): | |
# sample['jpg'] is tensor hwc in [-1, 1] at this point | |
x = sample[self.image_key] | |
w = torch.rand(3, device=x.device) | |
w /= w.sum() | |
out = torch.einsum('hwc,c->hw', x, w) | |
# Keep as 3ch so we can pass to encoder, also we might want to add hints | |
sample['lr'] = out.unsqueeze(-1).tile(1,1,3) | |
return sample | |
class AddMask(PRNGMixin): | |
def __init__(self, mode="512train", p_drop=0.): | |
super().__init__() | |
assert mode in list(MASK_MODES.keys()), f'unknown mask generation mode "{mode}"' | |
self.make_mask = MASK_MODES[mode] | |
self.p_drop = p_drop | |
def __call__(self, sample): | |
# sample['jpg'] is tensor hwc in [-1, 1] at this point | |
x = sample['jpg'] | |
mask = self.make_mask(self.prng, x.shape[0], x.shape[1]) | |
if self.prng.choice(2, p=[1 - self.p_drop, self.p_drop]): | |
mask = np.ones_like(mask) | |
mask[mask < 0.5] = 0 | |
mask[mask > 0.5] = 1 | |
mask = torch.from_numpy(mask[..., None]) | |
sample['mask'] = mask | |
sample['masked_image'] = x * (mask < 0.5) | |
return sample | |
class AddEdge(PRNGMixin): | |
def __init__(self, mode="512train", mask_edges=True): | |
super().__init__() | |
assert mode in list(MASK_MODES.keys()), f'unknown mask generation mode "{mode}"' | |
self.make_mask = MASK_MODES[mode] | |
self.n_down_choices = [0] | |
self.sigma_choices = [1, 2] | |
self.mask_edges = mask_edges | |
def __call__(self, sample): | |
# sample['jpg'] is tensor hwc in [-1, 1] at this point | |
x = sample['jpg'] | |
mask = self.make_mask(self.prng, x.shape[0], x.shape[1]) | |
mask[mask < 0.5] = 0 | |
mask[mask > 0.5] = 1 | |
mask = torch.from_numpy(mask[..., None]) | |
sample['mask'] = mask | |
n_down_idx = self.prng.choice(len(self.n_down_choices)) | |
sigma_idx = self.prng.choice(len(self.sigma_choices)) | |
n_choices = len(self.n_down_choices)*len(self.sigma_choices) | |
raveled_idx = np.ravel_multi_index((n_down_idx, sigma_idx), | |
(len(self.n_down_choices), len(self.sigma_choices))) | |
normalized_idx = raveled_idx/max(1, n_choices-1) | |
n_down = self.n_down_choices[n_down_idx] | |
sigma = self.sigma_choices[sigma_idx] | |
kernel_size = 4*sigma+1 | |
kernel_size = (kernel_size, kernel_size) | |
sigma = (sigma, sigma) | |
canny = kornia.filters.Canny( | |
low_threshold=0.1, | |
high_threshold=0.2, | |
kernel_size=kernel_size, | |
sigma=sigma, | |
hysteresis=True, | |
) | |
y = (x+1.0)/2.0 # in 01 | |
y = y.unsqueeze(0).permute(0, 3, 1, 2).contiguous() | |
# down | |
for i_down in range(n_down): | |
size = min(y.shape[-2], y.shape[-1])//2 | |
y = kornia.geometry.transform.resize(y, size, antialias=True) | |
# edge | |
_, y = canny(y) | |
if n_down > 0: | |
size = x.shape[0], x.shape[1] | |
y = kornia.geometry.transform.resize(y, size, interpolation="nearest") | |
y = y.permute(0, 2, 3, 1)[0].expand(-1, -1, 3).contiguous() | |
y = y*2.0-1.0 | |
if self.mask_edges: | |
sample['masked_image'] = y * (mask < 0.5) | |
else: | |
sample['masked_image'] = y | |
sample['mask'] = torch.zeros_like(sample['mask']) | |
# concat normalized idx | |
sample['smoothing_strength'] = torch.ones_like(sample['mask'])*normalized_idx | |
return sample | |
def example00(): | |
url = "pipe:aws s3 cp s3://s-datasets/laion5b/laion2B-data/000000.tar -" | |
dataset = wds.WebDataset(url) | |
example = next(iter(dataset)) | |
for k in example: | |
print(k, type(example[k])) | |
print(example["__key__"]) | |
for k in ["json", "txt"]: | |
print(example[k].decode()) | |
image = Image.open(io.BytesIO(example["jpg"])) | |
outdir = "tmp" | |
os.makedirs(outdir, exist_ok=True) | |
image.save(os.path.join(outdir, example["__key__"] + ".png")) | |
def load_example(example): | |
return { | |
"key": example["__key__"], | |
"image": Image.open(io.BytesIO(example["jpg"])), | |
"text": example["txt"].decode(), | |
} | |
for i, example in tqdm(enumerate(dataset)): | |
ex = load_example(example) | |
print(ex["image"].size, ex["text"]) | |
if i >= 100: | |
break | |
def example01(): | |
# the first laion shards contain ~10k examples each | |
url = "pipe:aws s3 cp s3://s-datasets/laion5b/laion2B-data/{000000..000002}.tar -" | |
batch_size = 3 | |
shuffle_buffer = 10000 | |
dset = wds.WebDataset( | |
url, | |
nodesplitter=wds.shardlists.split_by_node, | |
shardshuffle=True, | |
) | |
dset = (dset | |
.shuffle(shuffle_buffer, initial=shuffle_buffer) | |
.decode('pil', handler=warn_and_continue) | |
.batched(batch_size, partial=False, | |
collation_fn=dict_collation_fn) | |
) | |
num_workers = 2 | |
loader = wds.WebLoader(dset, batch_size=None, shuffle=False, num_workers=num_workers) | |
batch_sizes = list() | |
keys_per_epoch = list() | |
for epoch in range(5): | |
keys = list() | |
for batch in tqdm(loader): | |
batch_sizes.append(len(batch["__key__"])) | |
keys.append(batch["__key__"]) | |
for bs in batch_sizes: | |
assert bs==batch_size | |
print(f"{len(batch_sizes)} batches of size {batch_size}.") | |
batch_sizes = list() | |
keys_per_epoch.append(keys) | |
for i_batch in [0, 1, -1]: | |
print(f"Batch {i_batch} of epoch {epoch}:") | |
print(keys[i_batch]) | |
print("next epoch.") | |
def example02(): | |
from omegaconf import OmegaConf | |
from torch.utils.data.distributed import DistributedSampler | |
from torch.utils.data import IterableDataset | |
from torch.utils.data import DataLoader, RandomSampler, Sampler, SequentialSampler | |
from pytorch_lightning.trainer.supporters import CombinedLoader, CycleIterator | |
#config = OmegaConf.load("configs/stable-diffusion/txt2img-1p4B-multinode-clip-encoder-high-res-512.yaml") | |
#config = OmegaConf.load("configs/stable-diffusion/txt2img-upscale-clip-encoder-f16-1024.yaml") | |
config = OmegaConf.load("configs/stable-diffusion/txt2img-v2-clip-encoder-improved_aesthetics-256.yaml") | |
datamod = WebDataModuleFromConfig(**config["data"]["params"]) | |
dataloader = datamod.train_dataloader() | |
for batch in dataloader: | |
print(batch.keys()) | |
print(batch["jpg"].shape) | |
break | |
def example03(): | |
# improved aesthetics | |
tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{000000..060207}.tar -" | |
dataset = wds.WebDataset(tars) | |
def filter_keys(x): | |
try: | |
return ("jpg" in x) and ("txt" in x) | |
except Exception: | |
return False | |
def filter_size(x): | |
try: | |
return x['json']['original_width'] >= 512 and x['json']['original_height'] >= 512 | |
except Exception: | |
return False | |
def filter_watermark(x): | |
try: | |
return x['json']['pwatermark'] < 0.5 | |
except Exception: | |
return False | |
dataset = (dataset | |
.select(filter_keys) | |
.decode('pil', handler=wds.warn_and_continue)) | |
n_save = 20 | |
n_total = 0 | |
n_large = 0 | |
n_large_nowm = 0 | |
for i, example in enumerate(dataset): | |
n_total += 1 | |
if filter_size(example): | |
n_large += 1 | |
if filter_watermark(example): | |
n_large_nowm += 1 | |
if n_large_nowm < n_save+1: | |
image = example["jpg"] | |
image.save(os.path.join("tmp", f"{n_large_nowm-1:06}.png")) | |
if i%500 == 0: | |
print(i) | |
print(f"Large: {n_large}/{n_total} | {n_large/n_total*100:.2f}%") | |
if n_large > 0: | |
print(f"No Watermark: {n_large_nowm}/{n_large} | {n_large_nowm/n_large*100:.2f}%") | |
def example04(): | |
# improved aesthetics | |
for i_shard in range(60208)[::-1]: | |
print(i_shard) | |
tars = "pipe:aws s3 cp s3://s-laion/improved-aesthetics-laion-2B-en-subsets/aesthetics_tars/{:06}.tar -".format(i_shard) | |
dataset = wds.WebDataset(tars) | |
def filter_keys(x): | |
try: | |
return ("jpg" in x) and ("txt" in x) | |
except Exception: | |
return False | |
def filter_size(x): | |
try: | |
return x['json']['original_width'] >= 512 and x['json']['original_height'] >= 512 | |
except Exception: | |
return False | |
dataset = (dataset | |
.select(filter_keys) | |
.decode('pil', handler=wds.warn_and_continue)) | |
try: | |
example = next(iter(dataset)) | |
except Exception: | |
print(f"Error @ {i_shard}") | |
if __name__ == "__main__": | |
#example01() | |
#example02() | |
example03() | |
#example04() | |