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