File size: 4,121 Bytes
9b91042 | 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 | """build_cache — images + captions -> the trainer's cache dict.
The trainer is cache-fed (the pod-bed discipline): latents and text
features are encoded ONCE, then training never touches the VAE/TE again.
The campaign's caches load directly too — any dict with lat/ehs
(+ optional blob, val_lat/val_ehs) works.
Blob targets: the campaign rasterized entity polygons from the fused
dataset (research record, exp011). The generic path here accepts a
directory of binary masks (PNG, same stem as each image), resized to the
latent grid. Empty masks are counted and reported LOUDLY (the exp011
schema lesson: an empty blob target is a silent zero).
"""
from __future__ import annotations
from pathlib import Path
import torch
@torch.no_grad()
def build_cache(image_dir, pipe, *, captions=None, n_val: int = 256,
size: int = 512, blob_masks=None, device: str = "cuda",
batch: int = 8) -> dict:
"""image_dir: folder of images. pipe: a diffusers pipeline carrying
vae/text_encoder/tokenizer. captions: None | txt file (one caption per
image, sorted order) | dict stem->caption. blob_masks: None | folder
of binary PNG masks (same stems)."""
from PIL import Image
import numpy as np
paths = sorted(p for p in Path(image_dir).iterdir()
if p.suffix.lower() in (".png", ".jpg", ".jpeg", ".webp"))
assert paths, f"no images under {image_dir}"
if isinstance(captions, (str, Path)):
lines = Path(captions).read_text(encoding="utf-8").splitlines()
assert len(lines) >= len(paths), \
f"{len(lines)} captions for {len(paths)} images"
cap_of = {p.stem: lines[i] for i, p in enumerate(paths)}
elif isinstance(captions, dict):
cap_of = captions
else:
cap_of = {p.stem: "" for p in paths}
vae, te, tok = pipe.vae, pipe.text_encoder, pipe.tokenizer
sf = vae.config.scaling_factor
lat_sz = size // (2 ** (len(vae.config.block_out_channels) - 1))
lats, ehss, blobs, empty = [], [], [], 0
for i in range(0, len(paths), batch):
chunk = paths[i:i + batch]
imgs = []
for p in chunk:
im = Image.open(p).convert("RGB").resize((size, size))
imgs.append(torch.from_numpy(
np.asarray(im)).float().permute(2, 0, 1) / 127.5 - 1)
x = torch.stack(imgs).to(device)
lats.append((vae.encode(x).latent_dist.sample(
generator=torch.Generator(device=device).manual_seed(1400 + i))
* sf).cpu())
t = tok([cap_of.get(p.stem, "") for p in chunk],
padding="max_length", truncation=True,
max_length=tok.model_max_length, return_tensors="pt")
ehss.append(te(t.input_ids.to(device))[0].cpu())
if blob_masks is not None:
for p in chunk:
mp = Path(blob_masks) / f"{p.stem}.png"
if mp.exists():
m = Image.open(mp).convert("L").resize(
(lat_sz, lat_sz), Image.NEAREST)
mm = (torch.from_numpy(np.asarray(m)) > 127).float()
else:
mm = torch.zeros(lat_sz, lat_sz)
if mm.sum() == 0:
empty += 1
blobs.append(mm)
lat = torch.cat(lats)
ehs = torch.cat(ehss)
cache = {"lat": lat[:-n_val] if n_val else lat,
"ehs": ehs[:-n_val] if n_val else ehs}
if n_val:
cache["val_lat"] = lat[-n_val:]
cache["val_ehs"] = ehs[-n_val:]
if blob_masks is not None:
bl = torch.stack(blobs)
cache["blob"] = bl[:-n_val] if n_val else bl
if n_val:
cache["val_blob"] = bl[-n_val:]
cache["blob_empty_count"] = empty
if empty:
print(f"WARNING: {empty}/{len(paths)} blob masks are EMPTY — "
"empty targets train nothing and read as silent zeros "
"(the exp011 schema lesson)", flush=True)
print(f"cache: {cache['lat'].shape[0]} train rows"
+ (f" + {n_val} val" if n_val else ""), flush=True)
return cache
|