LaVIT

git clone https://github.com/jy0205/LaVIT.git
cd LaVIT
pip install -r requirements.txt

• (Optional) We recommend to use memory efficient attention by installing xFormers following the instructions in here. Then, you can set the argument use_xformers=True in build_model function to save the GPU memory and speed up inference.

Model Zoo

We release the LaVIT weight that is built upon Llama-2-7B as the large language model.

Note: Due to the license restrictions of Llama1, we cannot publish its weights. Thus, we release the weight of LaVIT based on the Llama2.

The pre-trained weight of LaVIT can be found on the huggingface from here, which will take around 22GB of disk space. LaVIT achieves state-of-the-arts performance on various multi-modal downstream tasks. The detailed quantitive results are shown as follows:

Zero-shot Multi-modal Understanding

Zero-shot Text-to-Image Generation

Usage

LaVIT can serve as a multi-modal generalist to perform both multi-modal comprehension and generation. Below, we provide some examples. Only a few lines of code are needed to use LaVIT for inference. We also provide the detailed examples in the following jupyter notebooks for learning how to interact with LaVIT.

• understanding.ipynb : examples for multi-modal understanding
• text2image_synthesis.ipynb: examples for the text-to-image generation.
• multimodal_synthesis.ipynb: examples for image synthesis with multi-modal prompts.

Multi-modal Understanding

import os
import random
import torch
import torch.nn as nn
from models import build_model
from PIL import Image

seed = 1234
random.seed(seed)
torch.manual_seed(seed)

# The local directory you save the LaVIT pre-trained weight,
# it will automatically download the checkpoint from huggingface
model_path = '/path/LaVIT_weight'

# Using BFloat16 during inference
model_dtype = 'bf16'  # Or set to fp16 to enable float16 inference

# Inference using GPU-0
device_id = 0
torch.cuda.set_device(device_id)
device = torch.device('cuda')

# Building LaVIT for understanding and load its weight from huggingface
model = build_model(model_path=model_path, model_dtype=model_dtype,
            device_id=device_id, use_xformers=False, understanding=True)
model = model.to(device)

# Image Captioning
image_path = 'demo/caption_image.jpg'
caption = model.generate({"image": image_path})[0]
print(caption)
# an old photo of a horse and buggy in front of a building

# Visual Question Answering
image_path = 'demo/qa_image.jpg'
question = "What's that drink in the glass?"
answer = model.predict_answers({"image": image_path, "text_input": question}, max_len=10)[0]
print("The answer is: ", answer)
# The answer is: orange juice

Text-to-Image Synthesis

For the Image generation, the Classifier-Free Guidance scale is important. A larger scale will encourage the model to generate samples highly related to the input prompt while sacrificing the image quality. We set guidance_scale_for_llm=4.0 by default, you can increase this scale (e.g., 5.0 or 6.0) to encourage the generated image to follow the semantics of given prompts. Besides, you can modify the ratio to enable to generate the images with different aspect ratios.

import os
import torch
import random
import torch.nn as nn
from models import build_model
from PIL import Image

seed = 1234
random.seed(seed)
torch.manual_seed(seed)

# The local directory you save the LaVIT pre-trained weight,
# it will automatically download the checkpoint from huggingface
model_path = '/path/LaVIT_weight'

# Using BFloat16 during inference
model_dtype = 'bf16'    # Or set to fp16 to enable float16 inference

# Inference using GPU-0
device_id = 0
torch.cuda.set_device(device_id)
device = torch.device('cuda')
torch_dtype = torch.bfloat16 if model_dtype=="bf16" else torch.float16

# Building LaVIT for Generation and load the weight from huggingface
# You can set `use_xformers=True` if have installed xformers to save GPU mempry and speed up
model = build_model(model_path=model_path, model_dtype=model_dtype, device_id=device_id,
       use_xformers=False, understanding=False, load_tokenizer=False)
model = model.to(device)

# Text-to-Image Generation
prompt = "a sculpture of a duck made of wool"

# LaVIT support 6 different image aspect ratios
ratio_dict = {
    '1:1' : (1024, 1024),
    '4:3' : (896, 1152),
    '3:2' : (832, 1216),
    '16:9' : (768, 1344),
    '2:3' : (1216, 832),
    '3:4' : (1152, 896),
}

# The image aspect ratio you want to generate
ratio = '1:1'
height, width = ratio_dict[ratio]

with torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
    images = model.generate_image(prompt, width=width, height=height,
    num_return_images=1, guidance_scale_for_llm=4.0, num_inference_steps=50)

images[0].save("output/i2t_output.jpg")

Evaluation

The batch evaluation code with multiple GPUs on the ado