rain1011/LaVIT-7B-v2

LaVIT: Unified Language-Vision Pretraining in LLM with Dynamic Discrete Visual Tokenization

This is the latest version (LaVITv2) for the multi-modal large language model: LaVIT. The inference code of LaVIT can be found in here.

In this version, We further improve LaVIT's image generation capability. In the updated version, 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.

[arXiv] [BibTeX]

Setup

Requirements

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.

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 latest pre-trained weight of LaVIT can be found on the huggingface from here, which will take around 25GB of disk space. We strongly recommend you to download and use the latest version of LaVIT. 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

Model	Image Captioning			Visual Question Answering
	COCO	NoCaps	Flickr30K	VQAv2	OK-VQA	GQA	VizWiz
Flamingo-3B	73.0	-	60.6	49.2	41.2	-	28.9
Flamingo-9B	79.4	-	61.5	51.8	44.7	-	28.8
OpenFlamingo-9B	79.5	-	59.5	52.7	37.8	-	27.5
MetaLM	82.2	-	43.4	41.1	11.4	-	-
Kosmos-1	84.7	-	67.1	51.0	-	-	29.2
Kosmos-2	-	-	80.5	51.1	-	-	-
BLIP-2 (Vicuna-7B)	-	107.5	74.9	-	-	41.3	25.3
BLIP-2 (Vicuna-13B)	-	103.9	71.6	-	-	32.3	19.6
CM3Leon-7B	61.6	-	-	47.6	-	-	37.6
Emu (LLaMA-1-13B)	112.4	-	-	52.0	38.2	-	34.2
LaVIT (LLaMA-1-7B)	134.0	114.2	83.0	66.0	54.6	46.8	38.5
LaVIT (LLaMA-2-7B)	134.6	113.1	83.2	68.2	55.7	48.0	45.3

Zero-shot Text-to-Image Generation

Method	Model	Model type	FID
Text2Image Specialist	DALL-E	Autoregressive	28.0
	CogView	Autoregressive	27.1
	StableDiffusion	Diffusion	12.6
	GLIDE	Diffusion	12.2
	DALL-E 2	Diffusion	10.4
	Make-A-Scene	Autoregressive	11.8
	MUSE-7.6B	Non-Autoregressive	7.9
	Imagen-3.4B	Diffusion	7.3
	Parti-20B	Autoregressive	7.2
Multimodal Large Langauge Model	GILL (OPT-6.7B)	LLM	12.2
	Emu (LLaMA-1-13B)	LLM	11.7
	CM3Leon-7B	LLM	10.8
	LaVIT (LLaMA-1-7B)	LLM	7.4
	LaVIT (LLaMA-2-7B)	LLM	7.2

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

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

Evaluation

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

Acknowledgement

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.

Citation

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}
}