File size: 4,045 Bytes
a012ee3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import os, sys, io, gc, time, tarfile, argparse, numpy as np, torch
from concurrent.futures import ThreadPoolExecutor
from PIL import Image
from huggingface_hub import hf_hub_download
from diffusers import AutoencoderKL
from transformers import CLIPTokenizer
REPO = 'undefined443/cc12m-wds-coco-recaptioned'
OUT = '/root/v5data'
TMP = '/root/v5tmp'
SCALE = 0.18215
RES = 256
MAXTOK = 40
ap = argparse.ArgumentParser()
ap.add_argument('--start', type=int, default=0)
ap.add_argument('--end', type=int, default=598)
ap.add_argument('--batch', type=int, default=64)
args = ap.parse_args()
os.makedirs(OUT, exist_ok=True)
os.makedirs(TMP, exist_ok=True)
dev = 'cuda'
vae = AutoencoderKL.from_pretrained('stabilityai/sd-vae-ft-mse').to(dev).half().eval()
for p in vae.parameters():
p.requires_grad_(False)
tok = CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32')
def prep_img(raw):
try:
im = Image.open(io.BytesIO(raw)).convert('RGB')
except Exception:
return None
w, h = im.size
s = RES / min(w, h)
im = im.resize((max(RES, int(w * s + 0.5)), max(RES, int(h * s + 0.5))), Image.BICUBIC)
w, h = im.size
l, t = (w - RES) // 2, (h - RES) // 2
return np.asarray(im.crop((l, t, l + RES, t + RES)), dtype=np.uint8)
pool = ThreadPoolExecutor(16)
t_start = time.time()
done_total = 0
def fetch(si):
return hf_hub_download(REPO, f'cc12m-coco-{si:04d}.tar', repo_type='dataset', local_dir=TMP)
todo = [i for i in range(args.start, args.end) if not os.path.exists(f'{OUT}/shard_{i:04d}.npz')]
print(f'{len(todo)} shards to process', flush=True)
dl = ThreadPoolExecutor(2)
nxt = dl.submit(fetch, todo[0]) if todo else None
for k, si in enumerate(todo):
dst = f'{OUT}/shard_{si:04d}.npz'
try:
p = nxt.result()
except Exception as e:
print(f' shard {si} download failed: {type(e).__name__}', flush=True)
nxt = dl.submit(fetch, todo[k + 1]) if k + 1 < len(todo) else None
continue
nxt = dl.submit(fetch, todo[k + 1]) if k + 1 < len(todo) else None
pairs = {}
try:
with tarfile.open(p) as tf:
for m in tf:
if not m.isfile():
continue
key, _, ext = m.name.partition('.')
if ext == 'jpg':
pairs.setdefault(key, {})['img'] = tf.extractfile(m).read()
elif ext == 'txt':
pairs.setdefault(key, {})['cap'] = tf.extractfile(m).read().decode('utf-8', 'ignore').strip()
except Exception as e:
print(f' shard {si} tar failed: {type(e).__name__}', flush=True)
os.remove(p); continue
os.remove(p)
items = [(v['img'], v['cap']) for v in pairs.values() if 'img' in v and 'cap' in v and v['cap']]
del pairs
if not items:
continue
lat_all, tok_all = [], []
for i in range(0, len(items), args.batch):
chunk = items[i:i + args.batch]
arrs = list(pool.map(prep_img, [c[0] for c in chunk]))
keep = [j for j, a in enumerate(arrs) if a is not None]
if not keep:
continue
x = np.stack([arrs[j] for j in keep])
t = torch.from_numpy(x).to(dev).half().permute(0, 3, 1, 2).div_(127.5).sub_(1.0)
with torch.no_grad():
lat = vae.encode(t).latent_dist.sample() * SCALE
lat_all.append(lat.cpu().numpy().astype(np.float16))
ids = tok([chunk[j][1] for j in keep], padding='max_length', max_length=MAXTOK,
truncation=True, return_tensors='np')['input_ids']
tok_all.append(ids.astype(np.int32))
if not lat_all:
continue
lat = np.concatenate(lat_all); ids = np.concatenate(tok_all)
np.savez(dst, latents=lat, tokens=ids)
done_total += len(lat)
el = time.time() - t_start
print(f' shard {si:04d}: {len(lat)} imgs total {done_total} {done_total/el:.0f} img/s '
f'{el/60:.1f} min elapsed', flush=True)
del lat_all, tok_all, lat, ids, items
gc.collect()
print('PREPDONE', done_total, flush=True)
|