chat.py

#!/usr/bin/env python3
"""AtlasVision inference: SigLIP2 vision encoder + N-ATLaS (Llama-3 8B) via a trained MLP projector.

  pip install -U torch transformers peft safetensors pillow accelerate huggingface_hub
  # optional for 4-bit on small GPUs (e.g. Colab T4):  pip install bitsandbytes
  export HF_TOKEN=hf_...   # needs accepted access to the gated NCAIR1/N-ATLaS

  python chat.py --image photo.jpg --question "What is happening in this picture?"          # stage 2
  python chat.py --stage 1 --image photo.jpg                                                  # stage 1 captioner
  python chat.py --image https://example.com/cat.jpg --question "Kedu ihe dị na foto a?"
  python chat.py --image photo.jpg --load-in-4bit                                             # ~8 GB GPU
  python chat.py --image photo.jpg --interactive                                              # several questions

Needs ~18 GB of GPU memory in bf16 (A100, L4, RTX 4090...); --load-in-4bit fits ~8 GB GPUs such as a Colab T4.
"""
import argparse
import io
import os
import sys

import torch
import torch.nn as nn
from PIL import Image

REPO = "FUTO-NIGERIA/AtlasVision"
LLM = "NCAIR1/N-ATLaS"
VISION = "google/siglip2-base-patch16-224"
USER_HEADER = "<|start_header_id|>user<|end_header_id|>\n\n"
ASSIST_HEADER = "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
EOT = "<|eot_id|>"


class ProjectionMLP(nn.Module):
    def __init__(self, vision_dim, text_dim):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(vision_dim, text_dim), nn.GELU(), nn.Linear(text_dim, text_dim))

    def forward(self, x):
        return self.net(x)


def fetch(repo, filename):
    if os.path.isdir(repo):
        return os.path.join(repo, filename)
    from huggingface_hub import hf_hub_download
    return hf_hub_download(repo, filename)


def load_image(src):
    if src.startswith(("http://", "https://")):
        import urllib.request
        with urllib.request.urlopen(src) as r:
            return Image.open(io.BytesIO(r.read())).convert("RGB")
    return Image.open(src).convert("RGB")


class AtlasVision:
    def __init__(self, stage=2, repo=REPO, llm=LLM, vision=VISION, load_in_4bit=False, device=None):
        from safetensors.torch import load_file
        from transformers import AutoImageProcessor, AutoModel, AutoModelForCausalLM, AutoTokenizer

        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        cuda = self.device.type == "cuda"
        self.dtype = torch.bfloat16 if (not cuda or torch.cuda.is_bf16_supported()) else torch.float16

        self.tok = AutoTokenizer.from_pretrained(llm)
        self.pad_id = self.tok.pad_token_id if self.tok.pad_token_id is not None else self.tok.eos_token_id
        self.eot_id = self.tok.convert_tokens_to_ids(EOT)
        self.processor = AutoImageProcessor.from_pretrained(vision)

        full_vision = AutoModel.from_pretrained(vision, dtype=self.dtype)
        self.vision = full_vision.vision_model.to(self.device).eval()
        vision_dim = full_vision.config.vision_config.hidden_size
        del full_vision

        kw = {"dtype": self.dtype}
        if load_in_4bit:
            from transformers import BitsAndBytesConfig
            kw["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                                                           bnb_4bit_compute_dtype=self.dtype)
            kw["device_map"] = {"": self.device.index or 0}
        self.llm = AutoModelForCausalLM.from_pretrained(llm, **kw)
        if not load_in_4bit:
            self.llm.to(self.device)
        text_dim = self.llm.config.hidden_size

        if stage == 2:
            from peft import PeftModel
            if os.path.isdir(repo):
                self.llm = PeftModel.from_pretrained(self.llm, os.path.join(repo, "stage2/lora_adapter"))
            else:
                self.llm = PeftModel.from_pretrained(self.llm, repo, subfolder="stage2/lora_adapter")
        self.llm.eval()

        self.projector = ProjectionMLP(vision_dim, text_dim)
        self.projector.load_state_dict(load_file(fetch(repo, f"stage{stage}/projector.safetensors")))
        self.projector.to(self.device, dtype=torch.float32).eval()

        self.prefix = torch.tensor([self.tok(USER_HEADER, add_special_tokens=True).input_ids], device=self.device)
        self.history = []  # (question, answer) turns about the current image

    @torch.no_grad()
    def ask(self, image, question, max_new_tokens=256, temperature=0.0):
        pv = self.processor(images=image, return_tensors="pt").pixel_values.to(self.device, self.dtype)
        img = self.projector(self.vision(pixel_values=pv).last_hidden_state.float()).to(self.dtype)

        text = ""
        for i, (q, a) in enumerate(self.history):
            text += (q if i == 0 else USER_HEADER + q) + ASSIST_HEADER + a + EOT
        text += (question if not self.history else USER_HEADER + question) + ASSIST_HEADER
        ids = torch.tensor([self.tok(text, add_special_tokens=False).input_ids], device=self.device)

        emb = self.llm.get_input_embeddings()
        embeds = torch.cat([emb(self.prefix).to(self.dtype), img, emb(ids).to(self.dtype)], dim=1)
        mask = torch.ones(embeds.shape[:2], dtype=torch.long, device=self.device)
        gen = dict(max_new_tokens=max_new_tokens, repetition_penalty=1.1, eos_token_id=self.eot_id, pad_token_id=self.pad_id)
        if temperature > 0:
            gen.update(do_sample=True, temperature=temperature, top_p=0.9)
        else:
            gen.update(do_sample=False)
        out = self.llm.generate(inputs_embeds=embeds, attention_mask=mask, **gen)
        answer = self.tok.decode(out[0], skip_special_tokens=True).strip()
        self.history.append((question, answer))
        return answer


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--image", required=True, help="path or URL")
    ap.add_argument("--question", default=None)
    ap.add_argument("--stage", type=int, default=2, choices=[1, 2])
    ap.add_argument("--repo", default=REPO, help="HF repo id, or a local folder containing stage1/ and stage2/")
    ap.add_argument("--llm", default=LLM)
    ap.add_argument("--vision", default=VISION)
    ap.add_argument("--load-in-4bit", action="store_true")
    ap.add_argument("--max-new-tokens", type=int, default=256)
    ap.add_argument("--temperature", type=float, default=0.0)
    ap.add_argument("--interactive", action="store_true")
    a = ap.parse_args()

    question = a.question or ("Describe this image briefly." if a.stage == 1 else "Describe this image in detail.")
    model = AtlasVision(a.stage, a.repo, a.llm, a.vision, a.load_in_4bit)
    image = load_image(a.image)
    print(f"\nQ: {question}\nA: {model.ask(image, question, a.max_new_tokens, a.temperature)}", flush=True)
    while a.interactive:
        try:
            q = input("\nQ (empty to quit): ").strip()
        except EOFError:
            break
        if not q:
            break
        print(f"A: {model.ask(image, q, a.max_new_tokens, a.temperature)}", flush=True)


if __name__ == "__main__":
    sys.exit(main())