File size: 2,774 Bytes
e3f7919 | 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 | """describe(image_path_or_feats, prompt) -> text, with the frozen Jais trunk + trained connector (stage 1: no LoRA expert).
image: a path / PIL image / bytes / list of those (encoded with encoder.py, SigLIP2 on CPU by default), or cached features (np/torch [576,768] or [B,576,768]).
CLI: python infer.py IMAGE [prompt]
env: CONNECTOR=dir holding connector.pt (default: the folder that contains this code/ folder);
DEVICE=cpu|cuda (SigLIP2 encoder device, default cpu); LOW=1 for the 64-token mode;
JAIS_PATH / SIGLIP2_PATH to use local base-model copies instead of downloading from the Hub."""
import io, os, sys, numpy as np, torch
HERE = os.path.dirname(os.path.abspath(__file__)); sys.path.insert(0, HERE)
from graft import Graft, load_trunk
_STATE = {}
def load_graft(connector_dir=None, device=None, trunk=None, tok=None):
"""one Graft per process; the connector file is 14 MB (7.08M params bf16 / 28 MB fp32 state dict)"""
# connector.pt sits next to the code/ folder in a downloaded release
connector_dir = connector_dir or os.environ.get("CONNECTOR") or os.path.dirname(HERE)
key = (connector_dir, device)
if key in _STATE: return _STATE[key]
device = device or ("cuda" if torch.cuda.is_available() else "cpu")
if trunk is None: tok, trunk = load_trunk(device)
g = Graft(trunk, tok); g.load_state(torch.load(os.path.join(connector_dir, "connector.pt"), map_location="cpu")); g.eval()
_STATE[key] = g; return g
def to_feats(image, enc_device="cpu"):
"""-> fp16/fp32 tensor [B,576,768]"""
if isinstance(image, torch.Tensor): f = image
elif isinstance(image, np.ndarray) and image.ndim >= 2 and image.shape[-2:] == (576, 768): f = torch.from_numpy(image)
else:
import encoder
from PIL import Image
items = image if isinstance(image, (list, tuple)) else [image]
items = [Image.open(io.BytesIO(x)) if isinstance(x, (bytes, bytearray)) else x for x in items]
f = torch.from_numpy(encoder.encode(items, device=enc_device))
return f.unsqueeze(0) if f.dim() == 2 else f
@torch.no_grad()
def describe(image, prompt="صِف هذه الصورة", graft=None, max_new=96, low=False, enc_device="cpu", **kw):
g = graft or load_graft(); f = to_feats(image, enc_device).to(g.markers.device)
qs = [prompt] * f.shape[0]
with torch.autocast("cuda" if g.markers.device.type == "cuda" else "cpu", dtype=torch.bfloat16):
out = g.generate(f, qs, max_new=max_new, low=low, **kw)
return out[0] if len(out) == 1 else out
if __name__ == "__main__":
img = sys.argv[1]; q = sys.argv[2] if len(sys.argv) > 2 else "صِف هذه الصورة"
print(describe(img, q, low=os.environ.get("LOW") == "1", enc_device=os.environ.get("DEVICE", "cpu")))
|