Views
No views yet
pip install torch transformers torchvision huggingface_hub tqdm matplotlib Pillow1import llm
2import anymodal
3import torch
4import vision
5from torch.utils.data import DataLoader
6import numpy as np
7import os
8from PIL import Image
9from huggingface_hub import hf_hub_download
10
11# Load language model and tokenizer
12llm_tokenizer, llm_model = llm.get_llm(
13 "meta-llama/Llama-3.2-1B",
14 access_token="GET_YOUR_OWN_TOKEN_FROM_HUGGINGFACE",
15 use_peft=False,
16)
17llm_hidden_size = llm.get_hidden_size(llm_tokenizer, llm_model)
18
19# Load vision model components
20image_processor, vision_model, vision_hidden_size = vision.get_image_encoder("google/vit-base-patch16-224", use_peft=False)
21
22# Initialize vision tokenizer and encoder
23vision_encoder = vision.VisionEncoder(vision_model)
24vision_tokenizer = vision.Projector(vision_hidden_size, llm_hidden_size, num_hidden=1)
25
26# Initialize MultiModalModel
27multimodal_model = anymodal.MultiModalModel(
28 input_processor=None,
29 input_encoder=vision_encoder,
30 input_tokenizer=vision_tokenizer,
31 language_tokenizer=llm_tokenizer,
32 language_model=llm_model,
33 input_start_token="<|imstart|>",
34 input_end_token="<|imend|>",
35 prompt_text="The description of the given image is: ",
36)
37
38# Download pre-trained model weights
39if not os.path.exists("image_captioning_model"):
40 os.makedirs("image_captioning_model")
41
42hf_hub_download("AnyModal/Image-Captioning-Llama-3.2-1B", filename="input_tokenizer.pt", local_dir="image_captioning_model")
43multimodal_model._load_model("image_captioning_model")
44
45# Generate caption for an image
46image_path = "example_image.jpg" # Path to your image
47image = Image.open(image_path).convert("RGB")
48processed_image = image_processor(image, return_tensors="pt")
49processed_image = {key: val.squeeze(0) for key, val in processed_image.items()} # Remove batch dimension
50
51# Generate caption
52generated_caption = multimodal_model.generate(processed_image, max_new_tokens=120)
53print("Generated Caption:", generated_caption)