This is the official repository for the multi-modal large language model: LaVIT.

[arXiv] [BibTeX]

• 2023.11.17 👏👏👏 We further improve LaVIT's image generation capability. In the updated version (LaVIT-v2), the aesthetic and prompt-alignment of generated images has been improved. The probability of watermark is also greatly reduced. The improvements are summarized as follows:

  • Using LaVIT to generate better synthetic captions for the noisy Laion-Aesthetic (Like DALL-E 3).
  • Add the high-aesthetic training images from the open-source JourneyDB dataset.
  • Using the 20M synthetic Laion-Aesthetic data and 4.2M JourneyDB data to further finetune the LLM for 8K steps.

• 2023.10.31 🌟🌟🌟 We update the high-resolution pixel decoder in LaVIT, which supports to generate high resolution (1024 * 1024 pixels), muliple aspect ratios (1:1, 4:3, 3:2, 16:9 ...) and high aesthetics images. The quality of generated images have been improved siginificantly.

• 2023.10.17 🚀🚀🚀 We release the pre-trained weight for LaVIT on the HuggingFace and provide the inference code of using it for both multi-modal understanding and generation.

Stay tuned for this repository! We will continue to optimize the performance of LaVIT and support more interesting and powerful capabilities in the future.

• The online demo for LaVIT
• Further improve the multi-modal image synthesis performance of LaVIT, including multi-modal image editing, subject-driven image synthesis ...
• Further improve the multi-modal understanding ability through instruction tuning.
• The finetuning code of LaVIT.
• Improve the image synthesis speed.

We propose LaVIT, a new general-purpose multi-modal foundation model that inherits the successful learning paradigm of LLM: predicting the next image / text token in an auto-regressive manner. LaVIT introduces a well-designed visual tokenizer to translate the non-linguistic image into a sequence of discrete tokens like a foreign language that LLM can read. Hence, both images and texts can be handled simultaneously under the unified generative objective. For more technical details, please refer to our paper.

After pre-training, LaVIT can serve as a multi-modal generalist to perform both multi-modal comprehension and generation without further fine-tuning. Specifically, it has the following capabilities:

• Read image contents, generate the captions and answer the questions.
• Text-to-image creation.
• Image synthesis via Multi-modal Prompt.

LaVIT can synthesis high quality (up to 1024 * 1024), muliple aspect ratios, and high aesthetics images based on the given image prompt, whose image generation ablility is competitive with state-of-the-art image generators (Parti, SDXL, and DALLE-3). The following are some examples.

Since both image and text are unified as discrete tokens in LLM, LaVIT can accept several modality combinations (e.g., text, image+text, image+image) as prompt to generate corresponding images without any fine-tuning.

Given the input images, LaVIT can read the image content and understand its semantics. For example, our model can give the caption for input image and answer the corresponding questions.

The code for this repo is tested with PyTorch 1.13.1 and CUDA 11.7. You should first install and configure the Pytorch Environment (including torch and torchvision) can then install the requirements with the following commands:

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

• (Optional) We recommend using 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.

We release the LaVIT weight that is built upon Llama-2-7B as the large language model. The latest pre-trained weight of LaVIT can be found on the huggingface, which will take around 26GB of disk space. We strongly recommend you to download and use the latest version: LaVIT-v2, since the image generation capability of this version has been greatly improved.

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

LaVIT achieves state-of-the-arts performance on various multi-modal downstream tasks. The detailed quantitive results are shown as follows:

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.

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

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=25)

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

The batch evaluation code with multiple GPUs on the adopted multi-modal benchmarks will be released in the following days.

We are grateful for the following awesome projects when implementing LaVIT:

• LLaMA: Open and Efficient Foundation Language Models
• BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models
• EVA-CLIP: Improved Training Techniques for CLIP at Scale
• BEIT: Masked Image Modeling with Vector-Quantized Visual Tokenizers
• Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch.

Consider giving this repository a star and cite LaVIT in your publications if it helps your research.

@article{jin2023unified,
  title={Unified Language-Vision Pretraining in LLM with Dynamic Discrete Visual Tokenization},
  author={Jin, Yang and Xu, Kun and Xu, Kun and Chen, Liwei and Liao, Chao and Tan, Jianchao and Mu, Yadong and others},
  journal={arXiv preprint arXiv:2309.04669},
  year={2023}
}