| """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_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"))) |
|
|
