File size: 2,158 Bytes
c793f45 | 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 | """Differentiable loader for InternViT (OpenGVLab/InternViT-300M-448px-V2_5)."""
from __future__ import annotations
import glob
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModel
from _shim import apply_pretrained_shims
REPO = "OpenGVLab/InternViT-300M-448px-V2_5"
IMAGE_SIZE = 448
MEAN = (0.485, 0.456, 0.406)
STD = (0.229, 0.224, 0.225)
_STUB = Path(__file__).resolve().parent / "vendored" / "internvit" / "flash_attention.py"
def _patch_flash_attention_files() -> None:
stub = _STUB.read_text()
hub_path = hf_hub_download(REPO, "flash_attention.py")
Path(hub_path).write_text(stub)
for path in glob.glob(
str(Path.home() / "workspace/hf-cache/modules/transformers_modules/**/flash_attention.py"),
recursive=True,
):
if "InternViT" in path or "InternViT_hyphen" in path:
Path(path).write_text(stub)
def load_internvit(dtype: torch.dtype = torch.bfloat16) -> nn.Module:
apply_pretrained_shims()
_patch_flash_attention_files()
config = AutoConfig.from_pretrained(REPO, trust_remote_code=True)
config.use_flash_attn = False
model = AutoModel.from_pretrained(
REPO,
config=config,
trust_remote_code=True,
dtype=dtype,
)
model.eval()
model.requires_grad_(False)
return model.to(device="cuda", dtype=dtype)
def image_feat_internvit(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
"""x: (1,3,H,W) in [0,1] on cuda, requires_grad=True. Returns (1, D) CLS feature."""
assert x.shape[0] == 1 and x.dim() == 4
mean = torch.tensor(MEAN, device=x.device, dtype=x.dtype).view(1, 3, 1, 1)
std = torch.tensor(STD, device=x.device, dtype=x.dtype).view(1, 3, 1, 1)
px = F.interpolate(x, size=(IMAGE_SIZE, IMAGE_SIZE), mode="bicubic", align_corners=False)
px = (px - mean) / std
dtype = next(model.parameters()).dtype
out = model(pixel_values=px.to(dtype=dtype))
feat = out.pooler_output.float()
if feat.dim() == 1:
feat = feat.unsqueeze(0)
return feat
|