turing-motors/Heron-NVILA-Lite-2B-hf

# We have developed this implementation with `transformers==4.57.1`. Other versions of transformers may also work, but we have not tested them.
pip install "transformers[torch]==4.57.1" pillow

from transformers import AutoConfig, AutoModelForImageTextToText, AutoProcessor

MODEL_NAME = "turing-motors/Heron-NVILA-Lite-2B-hf"

# you can use config
config = AutoConfig.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModelForImageTextToText.from_config(config, trust_remote_code=True)

# or directly from_pretrained
model = AutoModelForImageTextToText.from_pretrained(MODEL_NAME, trust_remote_code=True, device_map="auto")

# load processor
processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True)

# show chat_template
print(processor.tokenizer.chat_template)

def generate_content(content: str, images: list | None = None, **kwargs) -> str:
    conversation = [{"role": "user", "content": content}]
    text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
    encoding = processor(text=text, images=images, return_tensors="pt").to(model.device)
    output = model.generate(**encoding, **kwargs)
    return processor.decode(output[0, len(encoding["input_ids"][0]):], skip_special_tokens=True)

# examples generate with raw text
response = generate_content("こんにちは")
print(response)
print("---" * 40)

# examples generate with text + image
from PIL import Image
import requests
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
response = generate_content("<image>\n画像を説明してください。", images=[image])
print(response)
print("---" * 40)

# examples generate using generation_config
from PIL import Image
import requests
from transformers import GenerationConfig
generation_config = {
    "max_new_tokens": 512,
    "temperature": 0.5,
    "do_sample": True,
}
generation_config = GenerationConfig(**generation_config)
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
response = generate_content(
    "<image>\n画像を説明してください。",
    images=[image],
    generation_config=generation_config
)
print(response)
print("---" * 40)

# examples generate with text + image + text + image + text
from PIL import Image
import requests
url_list = [
    "https://images.unsplash.com/photo-1694831404826-3400c48c188d?q=80&w=2070&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D",
    "https://images.unsplash.com/photo-1693240876439-473af88b4ed7?q=80&w=1974&auto=format&fit=crop&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D"
]
images = [
   Image.open(requests.get(url, stream=True).raw).convert("RGB") for url in url_list
]
response = generate_content("".join([
    "<image>\n",
    "これは日本の画像です",
    "<image>\n",
    "これはオーストリアの画像です",
    "各画像の違いを説明して"]), images)
print(response)
print("---" * 40)