janus_multimodal_image_understanding_vqa_inference.py

1import torch
2from transformers import AutoModelForCausalLM
3from janus.models import MultiModalityCausalLM, JanusVLCPTProcessor
4from janus.utils.io import load_pil_images
5
6# 1. Specify the model path
7model_path = "deepseek-ai/Janus-1.3B"
8
9# 2. Load the processor and model
10vl_chat_processor: JanusVLCPTProcessor = JanusVLCPTProcessor.from_pretrained(model_path)
11tokenizer = vl_chat_processor.tokenizer
12
13vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
14    model_path, trust_remote_code=True
15)
16vl_gpt = vl_gpt.to(torch.bfloat16).cuda().eval()
17
18# 3. Prepare the conversation input
19# Use '<image_placeholder>' to indicate where the image is located in the prompt
20conversation = [
21    {
22        "role": "User",
23        "content": "<image_placeholder>\nDescribe this image in detail.",
24        "images": ["images/dog.jpg"],
25    },
26    {"role": "Assistant", "content": ""},
27]
28
29# 4. Process the inputs
30pil_images = load_pil_images(conversation)
31prepare_inputs = vl_chat_processor(
32    conversations=conversation, images=pil_images, force_batchify=True
33).to(vl_gpt.device)
34
35# 5. Run inference
36inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
37
38outputs = vl_gpt.generate(
39    inputs_embeds=inputs_embeds,
40    attention_mask=prepare_inputs.attention_mask,
41    pad_token_id=tokenizer.eos_token_id,
42    bos_token_id=tokenizer.bos_token_id,
43    eos_token_id=tokenizer.eos_token_id,
44    max_new_tokens=512,
45    do_sample=False,
46    use_cache=True,
47)
48
49# 6. Decode and print the result
50answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
51print(f"{prepare_inputs['sft_format'][0]}", answer)