| """PathGen-CLIP-L (jamessyx/PathGen-CLIP-L, CLIP trained on PathGen-1.6M) text + image towers. |
| Usage: python 16_extract_pathgenclip.py {text|image} [--pathgen|--prompts] |
| text -> cache/text_pathgenclip.npy (N, 4, 768) ORION templated texts (cache/text.json) |
| text --pathgen -> cache/pathgen_text_pathgenclip.npy (N, 1, 768) PathGen captions (cache/pathgen_meta.csv) |
| text --prompts -> cache/prompts_text_pathgenclip.npy (P, 1, 768) concept prompts (prompts.json) |
| image -> cache/he_pathgenclip.npy (N, 8, 768) ORION tiles, 8 dihedral views |
| image --pathgen -> cache/pathgen_he_pathgenclip.npy (N, 8, 768) PathGen images, 8 dihedral views |
| All float16, projected CLIP embeddings (encode_text / encode_image), NOT L2-normalised -- same as 04b / 08. |
| |
| Architecture: the README's 'ViT-B-16' is wrong for -L. The checkpoint has visual.conv1 (1024,3,14,14), 24 visual |
| blocks, visual.positional_embedding (577,1024) = 24x24+1 -> 336px, text width 768 / 12 blocks, proj -> 768. |
| That is OpenAI ViT-L/14@336; the authors' HF conversion (jamessyx/pathgenclip-vit-large-patch14-hf) is |
| initialised from openai/clip-vit-large-patch14-336 with hidden_act=quick_gelu -> open_clip 'ViT-L-14-336-quickgelu', |
| loaded strictly. Preprocess: 336px bicubic, OpenAI-CLIP mean/std (all ORION / PathGen images are square, so the |
| open_clip Resize+CenterCrop equals the plain square resize used in 04b). Tokenizer: OpenAI CLIP BPE, 77 tokens. |
| """ |
| import json |
| import os |
| import sys |
|
|
| import numpy as np |
| import torch |
| from PIL import Image, PngImagePlugin |
|
|
| PngImagePlugin.MAX_TEXT_CHUNK = 64 * 1024 * 1024 |
|
|
| import open_clip |
| from huggingface_hub import hf_hub_download |
|
|
| from common import CACHE, ORION, ROOT, dihedral, load_all_tiles |
|
|
| ARCH = "ViT-L-14-336-quickgelu" |
| NAME = "pathgenclip" |
| mode = sys.argv[1] |
| SET = sys.argv[sys.argv.index("--set") + 1] if "--set" in sys.argv else "pathgen" if "--pathgen" in sys.argv else None |
| PATHGEN = SET is not None |
| PROMPTS = "--prompts" in sys.argv |
| dev = "cuda" |
| CLIPN = ([0.48145466, 0.4578275, 0.40821073], [0.26862954, 0.26130258, 0.27577711]) |
|
|
|
|
| def load_model(): |
| ckpt = hf_hub_download("jamessyx/PathGen-CLIP-L", "pathgen-clip-l.pt") |
| model = open_clip.create_model(ARCH, pretrained=None) |
| sd = torch.load(ckpt, map_location="cpu") |
| sd = sd.get("state_dict", sd) |
| sd = {k[len("module."):] if k.startswith("module.") else k: v for k, v in sd.items()} |
| model.load_state_dict(sd, strict=True) |
| return model.eval().to(dev) |
|
|
|
|
| model = load_model() |
|
|
| if mode == "text": |
| tok = open_clip.get_tokenizer(ARCH) |
| if PROMPTS: |
| pr = json.load(open(os.path.join(ROOT, "prompts.json"))) |
| flat = [v["text"] for v in pr.values()] |
| ids, NV = list(range(len(flat))), 1 |
| elif PATHGEN: |
| import pandas as pd |
| flat = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).caption.tolist() |
| ids, NV = list(range(len(flat))), 1 |
| else: |
| texts = json.load(open(os.path.join(CACHE, "text.json"))) |
| ids, NV = sorted(int(k) for k in texts), 4 |
| flat = [texts[str(i)][v] for i in ids for v in range(NV)] |
| out = [] |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16): |
| for s in range(0, len(flat), 256): |
| out.append(model.encode_text(tok(flat[s:s + 256]).to(dev)).float().cpu().numpy()) |
| out = np.concatenate(out).reshape(len(ids), NV, -1).astype(np.float16) |
| assert ids == list(range(len(ids))) |
| pre = "prompts_" if PROMPTS else f"{SET}_" if PATHGEN else "" |
| fn = os.path.join(CACHE, pre + f"text_{NAME}.npy") |
|
|
| elif mode == "image": |
| size = model.visual.image_size[0] if isinstance(model.visual.image_size, (tuple, list)) else model.visual.image_size |
| assert size == 336, size |
| mean = torch.tensor(CLIPN[0]).view(1, 3, 1, 1) |
| std = torch.tensor(CLIPN[1]).view(1, 3, 1, 1) |
| if PATHGEN: |
| import pandas as pd |
| paths = pd.read_csv(os.path.join(CACHE, f"{SET}_meta.csv")).image_path.values |
| else: |
| paths = np.array([os.path.join(ORION, p) for p in load_all_tiles().image_path]) |
| out = None |
| bs = 32 |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16): |
| for s in range(0, len(paths), bs): |
| ims = [np.asarray(Image.open(p).convert("RGB").resize((size, size), Image.BICUBIC)) |
| for p in paths[s:s + bs]] |
| x = (torch.from_numpy(np.stack(ims)).permute(0, 3, 1, 2).float() / 255 - mean) / std |
| for d in range(8): |
| e = model.encode_image(dihedral(x, d).to(dev)).float().cpu().numpy() |
| if out is None: |
| out = np.zeros((len(paths), 8, e.shape[1]), np.float16) |
| out[s:s + bs, d] = e |
| if s % 640 == 0: |
| print(NAME, s, len(paths), flush=True) |
| fn = os.path.join(CACHE, (f"{SET}_" if PATHGEN else "") + f"he_{NAME}.npy") |
| else: |
| raise ValueError(mode) |
|
|
| np.save(fn, out) |
| print("saved", fn, out.shape, out.dtype, "finite:", np.isfinite(out).all(), |
| "peak GPU MB:", torch.cuda.max_memory_allocated() // 2**20) |
|
|
