Upload via givemeanode export_data

README.md

@@ -15,7 +15,6 @@ language:
 datasets:
 - liuhaotian/LLaVA-Pretrain
 - liuhaotian/LLaVA-Instruct-150K
-- lmms-lab/POPE
 tags:
 - vision-language
 - multimodal
@@ -63,7 +62,7 @@ stage2/projector.safetensors          # stage-2 projector (use together with the
 stage2/lora_adapter/                  # PEFT LoRA adapter for NCAIR1/N-ATLaS
 stage2b/projector.safetensors         # RECOMMENDED: de-biased projector
 stage2b/lora_adapter/                 # RECOMMENDED: de-biased LoRA adapter
-chat.py                               # standalone inference script (CLI + Python API)
+chat.py                               # standalone inference script (CLI + Python API, incl. lang= cascade)
 code/                                 # exact training, evaluation and launch scripts used
 eval/                                 # raw evaluation results (JSON) and report
 logs/                                 # per-step training logs (loss, grad norm, LR, speed)
@@ -76,7 +75,7 @@ Accept the conditions on the [N-ATLaS model page](https://huggingface.co/NCAIR1/
 ```bash
 pip install -U torch transformers peft safetensors pillow accelerate huggingface_hub
 export HF_TOKEN=hf_...            # a token from the account that was granted N-ATLaS access
-wget https://huggingface.co/FUTO-NIGERIA/AtlasVision/resolve/main/chat.py
+wget https://huggingface.co/Modularcomputing/AtlasVision/resolve/main/chat.py
 
 python chat.py --image photo.jpg --question "What is happening in this picture?"   # add --stage 2b
 python chat.py --image photo.jpg --question "Kedu ihe dị na foto a?"      # Igbo
@@ -147,7 +146,7 @@ roughly doubles the loss — the language model is relying on the visual content
 
 | Metric | Stage 1 | Stage 2 | Stage 2b (recommended) |
 |---|---|---|---|
-| Held-out LLaVA-Instruct loss (500 unseen conversations, lower is better) | 2.2521 | 1.1271 | **1.1602** |
+| Held-out LLaVA-Instruct loss (500 unseen conversations, lower is better) | 2.2521 | **1.1271** | 1.1602 |
 
 **POPE** (object hallucination: yes/no questions about COCO val2014 images, scored from the Yes/No token
 probabilities; the benchmark is balanced 50% yes / 50% no, so a yes-ratio near 0.5 is ideal):
@@ -169,10 +168,17 @@ Stage 2b continues stage 2 for one epoch on a 150k mix drawn from the public LLa
 short-answer VQA conversations (**57,804 "yes" vs 59,443 "no"** turns) plus 40k LLaVA-Instruct conversations
 replayed so detailed description is not lost. Every POPE test image and every held-out conversation was removed
 from this mix before training, so the numbers above are not contaminated; region/bounding-box tasks were dropped
-as unsupported. Result: the yes-ratio fell to 0.56–0.63, accuracy rose by **24–28 points**, and the held-out
-conversation loss improved as well (1.1271 → 1.1602),
-so description quality did not regress. For reference, LLaVA-1.5-7B reports POPE F1 ≈ 0.86 from 665k samples at
-336 px; AtlasVision reaches 0.83 from 150k samples at 224 px.
+as unsupported. Result: the yes-ratio fell to 0.56–0.63 and accuracy rose by **24–28 points**.
+
+The held-out conversation loss rose slightly (1.1271 → 1.1602) — a small regression from the changed data mix.
+Description quality looks preserved in the examples below, but we have not measured that quantitatively.
+
+Averaged over the three splits, stage 2b scores F1 0.80. We deliberately do not put this head-to-head with
+the number quoted for LLaVA-1.5-7B: that paper reports a single POPE figure without stating how it aggregates the
+splits, it covers the full POPE benchmark (COCO + A-OKVQA + GQA) while we evaluate the COCO split only, and
+independent reproductions of LLaVA-1.5 on COCO with greedy decoding land anywhere from F1 ≈ 0.78 to ≈ 0.87
+depending on setup. Our setup also differs in scale (150k vs 665k samples), resolution (224 px vs 336 px) and base
+model, so treat these as our own baseline, not a ranking.
 
 Recall (0.88) still exceeds precision
 (0.70) on the adversarial split, so a mild "yes" lean remains on the
@@ -214,7 +220,7 @@ In addition to the vehicles, there are multiple people walking around the area.
 
 ### Questions in Nigerian languages (stage 2b)
 
-Stage 2b answers Igbo and Yoruba questions **in those languages** without any translation step (stage 2 always replied in English). Hausa still falls back to English, so the `lang=` cascade below remains the reliable route for all three.
+In our tests, stage 2b answered the Igbo and Yoruba questions **in those languages** without any translation step, where stage 2 always replied in English. This is **one example per language**, both short and containing English loanwords ("skier", "snow"), so treat it as an encouraging signal rather than a measured result — a proper multilingual evaluation is still to come. Hausa fell back to English. For reliable answers in all three languages, use the `lang=` cascade described below.
 
 | Language | Question | Answer |
 |---|---|---|
@@ -233,6 +239,36 @@ Does the stage-2 LoRA damage N-ATLaS's original text abilities? Same Igbo questi
 
 Both answer fluently in Igbo, so the adapter keeps N-ATLaS's text abilities intact (responses are cut at 120 tokens).
 
+## Answers in Igbo, Yoruba and Hausa (cascade)
+
+The image training data is English, so AtlasVision reasons about images in English. `chat.py` can route any of the
+four languages through **one** loaded model, switching the LoRA adapter off to recover the original N-ATLaS for the
+translation steps:
+
+1. **Question → English** with plain N-ATLaS (adapter off).
+2. **Look and answer in English** with AtlasVision (adapter on, image tokens in).
+3. **Answer → target language** with plain N-ATLaS (adapter off).
+
+```python
+from chat import AtlasVision, load_image
+
+model = AtlasVision(stage="2b")
+image = load_image("photo.jpg")
+
+r = model.chat(image, "Mutane nawa ne a cikin hoton?", lang="ha")
+print(r["answer"])            # Hausa
+print(r["english_answer"])    # English draft, so you can check what was translated
+
+print(model.describe(image, lang="ig"))        # Igbo description
+print(model.translate("Good morning", "yo"))   # plain N-ATLaS translation, no image
+```
+
+From the command line: `python chat.py --image photo.jpg --question "Kedu ihe di na foto a?" --lang ig`.
+
+Costs and caveats: each answer takes roughly 2–3× longer than English-only; translation quality is N-ATLaS's own; and
+any error in the English answer is carried into the translation, which is why `english_answer` is always returned.
+Pass `translate_question=False` to skip step 1 when the question is already short and simple.
+
 ## Limitations
 
 - **Hallucination.** Long descriptions are fluent but often add plausible details that are not in the image
@@ -257,7 +293,11 @@ This repository contains weights trained on top of other people's models and dat
   conditions on the [model page](https://huggingface.co/NCAIR1/N-ATLaS).
 - **SigLIP2**: Apache-2.0.
 - **LLaVA-Instruct-150K** (stage 2 data): CC BY-NC 4.0, generated with GPT-4 and subject to OpenAI's terms —
-  **stage-2 weights are for non-commercial research use.**
+  **stage-2 and stage-2b weights are for non-commercial research use.**
+- **LLaVA-1.5 mix / VQAv2, OK-VQA, A-OKVQA** (stage 2b data): drawn from `llava_v1_5_mix665k`, which inherits the
+  CC BY-NC 4.0 terms above. The underlying short-answer sets carry their own licenses (VQAv2 annotations are
+  CC BY 4.0 over COCO images; OK-VQA and A-OKVQA are research datasets released by their authors) — check each
+  dataset's page before any redistribution or commercial use.
 - **LLaVA-Pretrain** (stage 1 data): captions/images from LAION, Conceptual Captions and SBU under their respective terms.
 - **COCO images**: Flickr images under their individual Creative Commons licenses.
 
@@ -273,9 +313,9 @@ This repository contains weights trained on top of other people's models and dat
 ```bibtex
 @misc{atlasvision2026,
   title  = {AtlasVision: a vision-language extension of N-ATLaS},
-  author = {FUTO-NIGERIA},
+  author = {Modularcomputing},
   year   = {2026},
-  url    = {https://huggingface.co/FUTO-NIGERIA/AtlasVision}
+  url    = {https://huggingface.co/Modularcomputing/AtlasVision}
 }
 @inproceedings{liu2023llava,
   title     = {Visual Instruction Tuning},

chat.py

@@ -2,18 +2,19 @@
 """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
+  python chat.py --image photo.jpg --question "What is happening in this picture?"
+  python chat.py --image photo.jpg --question "Kedu ihe di na foto a?" --lang ig
+  python chat.py --image photo.jpg --stage 1          # stage-1 captioner
+  python chat.py --image photo.jpg --load-in-4bit     # ~8 GB GPU
+  python chat.py --image photo.jpg --interactive
 
-Needs ~18 GB of GPU memory in bf16 (A100, L4, RTX 4090...); --load-in-4bit fits ~8 GB GPUs such as a Colab T4.
+Stages: "2b" (default, de-biased), "2" (kept for reproducibility), "1" (captioner).
+Needs ~18 GB of GPU memory in bf16; --load-in-4bit fits ~8 GB GPUs such as a Colab T4.
 """
 import argparse
+import contextlib
 import io
 import os
 import sys
@@ -22,12 +23,13 @@ import torch
 import torch.nn as nn
 from PIL import Image
 
-REPO = "FUTO-NIGERIA/AtlasVision"
+REPO = "Modularcomputing/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|>"
+LANGUAGES = {"en": "English", "ig": "Igbo", "yo": "Yoruba", "ha": "Hausa"}
 
 
 class ProjectionMLP(nn.Module):
@@ -47,18 +49,25 @@ def fetch(repo, filename):
 
 
 def load_image(src):
-    if src.startswith(("http://", "https://")):
+    if isinstance(src, Image.Image):
+        return src.convert("RGB")
+    if isinstance(src, str) and 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")
 
 
+def lang_name(lang):
+    return "English" if lang is None else LANGUAGES.get(str(lang).lower(), str(lang).title())
+
+
 class AtlasVision:
-    def __init__(self, stage=2, repo=REPO, llm=LLM, vision=VISION, load_in_4bit=False, device=None):
+    def __init__(self, stage="2b", 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.stage = str(stage)
         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
@@ -84,24 +93,32 @@ class AtlasVision:
             self.llm.to(self.device)
         text_dim = self.llm.config.hidden_size
 
-        if stage == 2:
+        if self.stage != "1":
             from peft import PeftModel
+            sub = f"stage{self.stage}/lora_adapter"
             if os.path.isdir(repo):
-                self.llm = PeftModel.from_pretrained(self.llm, os.path.join(repo, "stage2/lora_adapter"))
+                self.llm = PeftModel.from_pretrained(self.llm, os.path.join(repo, sub))
             else:
-                self.llm = PeftModel.from_pretrained(self.llm, repo, subfolder="stage2/lora_adapter")
+                self.llm = PeftModel.from_pretrained(self.llm, repo, subfolder=sub)
         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.load_state_dict(load_file(fetch(repo, f"stage{self.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
+        self.history = []
+
+    def _gen(self, max_new_tokens, temperature, **inputs):
+        kw = dict(max_new_tokens=max_new_tokens, repetition_penalty=1.1,
+                  eos_token_id=self.eot_id, pad_token_id=self.pad_id)
+        kw.update(dict(do_sample=True, temperature=temperature, top_p=0.9) if temperature > 0 else dict(do_sample=False))
+        return self.llm.generate(**inputs, **kw)
 
     @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)
+        """One vision-language turn, in English. Follow-ups reuse self.history."""
+        pv = self.processor(images=load_image(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 = ""
@@ -113,23 +130,52 @@ class AtlasVision:
         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)
+        out = self._gen(max_new_tokens, temperature, inputs_embeds=embeds, attention_mask=mask)
         answer = self.tok.decode(out[0], skip_special_tokens=True).strip()
         self.history.append((question, answer))
         return answer
 
+    @torch.no_grad()
+    def text(self, prompt, max_new_tokens=512, temperature=0.0):
+        """Plain N-ATLaS: LoRA switched off, no image. Translation, explanation, any text task."""
+        ids = self.tok(USER_HEADER + prompt + ASSIST_HEADER, add_special_tokens=True,
+                       return_tensors="pt").input_ids.to(self.device)
+        off = self.llm.disable_adapter() if hasattr(self.llm, "disable_adapter") else contextlib.nullcontext()
+        with off:
+            out = self._gen(max_new_tokens, temperature, input_ids=ids, attention_mask=torch.ones_like(ids))
+        return self.tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True).strip()
+
+    def translate(self, text, target, source="English", max_new_tokens=512):
+        target, source = lang_name(target), lang_name(source)
+        if target == source:
+            return text
+        return self.text(f"Translate this {source} text to {target}. Reply with only the translation.\n\n{text}",
+                         max_new_tokens=max_new_tokens)
+
+    def chat(self, image, question, lang="en", translate_question=True, max_new_tokens=256, temperature=0.0):
+        """Ask about an image in English (en), Igbo (ig), Yoruba (yo) or Hausa (ha).
+
+        Cascade on one loaded model: question -> English (plain N-ATLaS) -> AtlasVision answers in English
+        -> answer -> target language (plain N-ATLaS). Returns both so the English draft can be checked.
+        """
+        name = lang_name(lang)
+        q_en = question if (name == "English" or not translate_question) else self.translate(question, "English", name)
+        a_en = self.ask(image, q_en, max_new_tokens, temperature)
+        answer = a_en if name == "English" else self.translate(a_en, name, "English", max_new_tokens=2 * max_new_tokens)
+        return {"answer": answer, "english_answer": a_en, "english_question": q_en}
+
+    def describe(self, image, lang="en", detailed=True, **kw):
+        q = "Describe this image in detail." if detailed else "Describe this image briefly."
+        return self.chat(image, q, lang=lang, translate_question=False, **kw)["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("--stage", default="2b", choices=["1", "2", "2b"])
+    ap.add_argument("--lang", default="en", choices=sorted(LANGUAGES))
+    ap.add_argument("--repo", default=REPO, help="HF repo id, or a local folder with stage1/ stage2/ stage2b/")
     ap.add_argument("--llm", default=LLM)
     ap.add_argument("--vision", default=VISION)
     ap.add_argument("--load-in-4bit", action="store_true")
@@ -138,10 +184,13 @@ def main():
     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.")
+    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)
+    r = model.chat(image, question, lang=a.lang, max_new_tokens=a.max_new_tokens, temperature=a.temperature)
+    print(f"\nQ: {question}\nA: {r['answer']}", flush=True)
+    if a.lang != "en":
+        print(f"[English draft: {r['english_answer']}]", flush=True)
     while a.interactive:
         try:
             q = input("\nQ (empty to quit): ").strip()
@@ -149,7 +198,8 @@ def main():
             break
         if not q:
             break
-        print(f"A: {model.ask(image, q, a.max_new_tokens, a.temperature)}", flush=True)
+        r = model.chat(image, q, lang=a.lang, max_new_tokens=a.max_new_tokens, temperature=a.temperature)
+        print(f"A: {r['answer']}", flush=True)
 
 
 if __name__ == "__main__":

code/chat.py

@@ -0,0 +1,206 @@
+#!/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
+  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?"
+  python chat.py --image photo.jpg --question "Kedu ihe di na foto a?" --lang ig
+  python chat.py --image photo.jpg --stage 1          # stage-1 captioner
+  python chat.py --image photo.jpg --load-in-4bit     # ~8 GB GPU
+  python chat.py --image photo.jpg --interactive
+
+Stages: "2b" (default, de-biased), "2" (kept for reproducibility), "1" (captioner).
+Needs ~18 GB of GPU memory in bf16; --load-in-4bit fits ~8 GB GPUs such as a Colab T4.
+"""
+import argparse
+import contextlib
+import io
+import os
+import sys
+
+import torch
+import torch.nn as nn
+from PIL import Image
+
+REPO = "Modularcomputing/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|>"
+LANGUAGES = {"en": "English", "ig": "Igbo", "yo": "Yoruba", "ha": "Hausa"}
+
+
+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 isinstance(src, Image.Image):
+        return src.convert("RGB")
+    if isinstance(src, str) and 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")
+
+
+def lang_name(lang):
+    return "English" if lang is None else LANGUAGES.get(str(lang).lower(), str(lang).title())
+
+
+class AtlasVision:
+    def __init__(self, stage="2b", 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.stage = str(stage)
+        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 self.stage != "1":
+            from peft import PeftModel
+            sub = f"stage{self.stage}/lora_adapter"
+            if os.path.isdir(repo):
+                self.llm = PeftModel.from_pretrained(self.llm, os.path.join(repo, sub))
+            else:
+                self.llm = PeftModel.from_pretrained(self.llm, repo, subfolder=sub)
+        self.llm.eval()
+
+        self.projector = ProjectionMLP(vision_dim, text_dim)
+        self.projector.load_state_dict(load_file(fetch(repo, f"stage{self.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 = []
+
+    def _gen(self, max_new_tokens, temperature, **inputs):
+        kw = dict(max_new_tokens=max_new_tokens, repetition_penalty=1.1,
+                  eos_token_id=self.eot_id, pad_token_id=self.pad_id)
+        kw.update(dict(do_sample=True, temperature=temperature, top_p=0.9) if temperature > 0 else dict(do_sample=False))
+        return self.llm.generate(**inputs, **kw)
+
+    @torch.no_grad()
+    def ask(self, image, question, max_new_tokens=256, temperature=0.0):
+        """One vision-language turn, in English. Follow-ups reuse self.history."""
+        pv = self.processor(images=load_image(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)
+        out = self._gen(max_new_tokens, temperature, inputs_embeds=embeds, attention_mask=mask)
+        answer = self.tok.decode(out[0], skip_special_tokens=True).strip()
+        self.history.append((question, answer))
+        return answer
+
+    @torch.no_grad()
+    def text(self, prompt, max_new_tokens=512, temperature=0.0):
+        """Plain N-ATLaS: LoRA switched off, no image. Translation, explanation, any text task."""
+        ids = self.tok(USER_HEADER + prompt + ASSIST_HEADER, add_special_tokens=True,
+                       return_tensors="pt").input_ids.to(self.device)
+        off = self.llm.disable_adapter() if hasattr(self.llm, "disable_adapter") else contextlib.nullcontext()
+        with off:
+            out = self._gen(max_new_tokens, temperature, input_ids=ids, attention_mask=torch.ones_like(ids))
+        return self.tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True).strip()
+
+    def translate(self, text, target, source="English", max_new_tokens=512):
+        target, source = lang_name(target), lang_name(source)
+        if target == source:
+            return text
+        return self.text(f"Translate this {source} text to {target}. Reply with only the translation.\n\n{text}",
+                         max_new_tokens=max_new_tokens)
+
+    def chat(self, image, question, lang="en", translate_question=True, max_new_tokens=256, temperature=0.0):
+        """Ask about an image in English (en), Igbo (ig), Yoruba (yo) or Hausa (ha).
+
+        Cascade on one loaded model: question -> English (plain N-ATLaS) -> AtlasVision answers in English
+        -> answer -> target language (plain N-ATLaS). Returns both so the English draft can be checked.
+        """
+        name = lang_name(lang)
+        q_en = question if (name == "English" or not translate_question) else self.translate(question, "English", name)
+        a_en = self.ask(image, q_en, max_new_tokens, temperature)
+        answer = a_en if name == "English" else self.translate(a_en, name, "English", max_new_tokens=2 * max_new_tokens)
+        return {"answer": answer, "english_answer": a_en, "english_question": q_en}
+
+    def describe(self, image, lang="en", detailed=True, **kw):
+        q = "Describe this image in detail." if detailed else "Describe this image briefly."
+        return self.chat(image, q, lang=lang, translate_question=False, **kw)["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", default="2b", choices=["1", "2", "2b"])
+    ap.add_argument("--lang", default="en", choices=sorted(LANGUAGES))
+    ap.add_argument("--repo", default=REPO, help="HF repo id, or a local folder with stage1/ stage2/ stage2b/")
+    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)
+    r = model.chat(image, question, lang=a.lang, max_new_tokens=a.max_new_tokens, temperature=a.temperature)
+    print(f"\nQ: {question}\nA: {r['answer']}", flush=True)
+    if a.lang != "en":
+        print(f"[English draft: {r['english_answer']}]", flush=True)
+    while a.interactive:
+        try:
+            q = input("\nQ (empty to quit): ").strip()
+        except EOFError:
+            break
+        if not q:
+            break
+        r = model.chat(image, q, lang=a.lang, max_new_tokens=a.max_new_tokens, temperature=a.temperature)
+        print(f"A: {r['answer']}", flush=True)
+
+
+if __name__ == "__main__":
+    sys.exit(main())