Skip to content

Latest commit

 

History

History
371 lines (320 loc) · 16.9 KB

README.md

File metadata and controls

371 lines (320 loc) · 16.9 KB

LANTERN

This repository is an official PyTorch implementation of the paper LANTERN: Accelerating Visual Autoregressive Models via Relaxed Speculative Decoding (ICLR 2025) and LANTERN++: Enhanced Relaxed Speculative Decoding with Static Tree Drafting for Visual Auto-regressive Models (ICLRW - SCOPE(Oral) 2025), which supports various functionalities related to LANTERN, including model inference, drafter model training, drafter model training data generation and image decoding for image generation.

📰 News

  • [2025-03-05] 🎉🎉🎉 LANTERN is released! 🎉🎉🎉

📚 Table of Contents

  1. Directory Structure
  2. Installation
  3. Key Features
  4. Usage
  5. License
  6. Acknowledgement
  7. Citation

🗂️ Directory Structure

The main directory structure of the project is as follows:

.
├── models/                                     # Model and related modules
│   ├── base_models/                            # Base model modules
│   │   ├── lumina_mgpt
│   │   │   ├── modeling_lumina_mgpt.py
│   │   │   └── other files...
│   │   └── other models...     
│   ├── kv_variants/                            # Key-Value variant models
│   │   ├── modeling_lumina_mgpt_kv.py
|   |   └── modeling_anole_kv.py
│   │   └── other models...
│   ├── drafters/                               # Drafter model modules
│   │   ├── kv_cache.py
│   │   ├── choices
│   │   ├── cnets_lumina_mgpt.py
|   |   ├── cnets_anole.py
│   │   ├── cnets_{other_models}.py ...   
│   │   └── utils.py
│   ├── configs/                                # Configuration modules
│   │   ├── configs.py
│   │   ├── configuration_lumina_mgpt.py
|   |   ├── configuration_anole.py
│   │   └── configuration_{other_models}.py...
│   ├── ea_model_lumina_mgpt.py                 # EAGLE models
|   ├── ea_model_anole.py
│   └── ea_model_{other_models}.py...
├── data/
│   ├── configs/
│   │   ├── lumina_mgpt_config.json             # Configuration for model init
|   |   ├── anole_config.json
│   │   └── configs for other models...
│   ├── prompts/                                # Prompts for image generation
│   ├── self_distilled_data/                    # Self-distilled data for drafter training
│   └── drafter_train_data/                     # Train data for drafter
├── ckpts/                                      # Model checkpoints folder
│   ├── lumina_mgpt/
│   │   ├── chameleon/                      
│   │   ├── Lumina-mGPT-7B-768/                 # Model and tokenizer files
│   │   ├── trained_drafters/                   # Trained drafter models
|   |   |   └──...state_20/
|   |   |      ├── config.json                  # config.json for drafter model
|   |   |      └── other files...
│   │   └── vq_distances/                       # Pre-computed VQ distances for LANTERN
│   └── other models...
├── entrypoints/                                # Execution entry points
│   ├── train_drafter/
│   │   ├── data_utils.py
│   │   └── main.py
│   ├── generate_codebook.py
│   ├── generate_images.py
│   ├── generate_train_data.py
│   └── other files...
├── third_party/                                # Third-party libraries
│   └── vllm
├── main.py                                     # Main execution script
├── requirements.txt                            # Project dependencies
├── environment.yaml
├── .gitignore
└── README.md

Here is a brief description for each directory.

  1. models/ - Contains model implementations and related modules.

    • base_models/ - Base model implementations (e.g., Lumina-mGPT, LlamaGen, Anole).
    • kv_variants/ - Modified base models with Key-Value cache adaptations for enhanced compatibility with EAGLE’s architecture.
    • drafters/ - Modules and auxiliary code for drafter models.
    • configs/ - Configuration modules for each model (e.g., ChameleonConfig for Lumina-mGPT).

  2. data/ - Stores configuration files, text prompts, self-distilled data, and drafter training data.

  3. ckpts/ - Checkpoints for all models, including trained drafters and VQ distances for relaxed speculative decoding.

  4. entrypoints/ - Primary scripts for tasks such as image generation, codebook generation, and drafter training.

  5. third_party/ - Custom external libraries, including modifications for specific functionality.

⚙️ Installation

  1. Install Required Packages Requirements

    • Python >= 3.10
    • PyTorch >= 2.4.0

    Install the dependencies listed in requirements.txt.

    git clone https://github.com/jadohu/LANTERN
cd LANTERN
pip install -r requirements.txt

  2. Additional Setup

    1. Lumina-mGPT For Lumina-mGPT, we need to install flash_attention and xllmx packages. 
      pip install flash-attn --no-build-isolation
cd models/base_models/lumina_mgpt
pip install -e .
      1. (Optional) vLLM Install and set up vLLM with the required modifications. Note that we use vLLM==0.6.3 and build from source. The required modifications are specifed in third_party/vllm. The installation procedure is as follows. 
        pip install https://vllm-wheels.s3.us-west-2.amazonaws.com/fd47e57f4b0d5f7920903490bce13bc9e49d8dba/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl
git clone https://github.com/vllm-project/vllm
cd vllm
git checkout tags/v0.6.3

cd ..
mv -rf third_party/vllm/* vllm/
cd vllm
python python_only_dev.py

  3. Checkpoints All model weights and other required data should be stored in ckpts/.

    1. Lumina-mGPT For Lumina-mGPT, since currently the Chameleon implementation in transformers does not contain the VQ-VAE decoder, please manually download the original VQ-VAE weights provided by Meta and put them to the following directory:

      ckpts
└── lumina_mgpt
    └── chameleon
        └── tokenizer
            ├── text_tokenizer.json
            ├── vqgan.yaml
            └── vqgan.ckpt

      Also download the original model Lumina-mGPT-7B-768 from Huggingface 🤗 and put them to the following directory:

      ckpts
└── lumina_mgpt
    └── Lumina-mGPT-7B-768
        ├── config.json
        ├── generation_config.json
        ├── model-00001-of-00002.safetensors
        └── other files...

    2. LlamaGen For LlamaGen T2I model, download LlamaGen-T2I and/or LlamaGen-T2I-2, which is a huggingface style converted model from LlamaGen.

      In addition, you should download VQ-VAE and flan-t5-xl.

      ckpts
└── llamagen
    ├── LlamaGen-T2I
    │   ├── config.json
    │   ├── generation_config.json
    │   ├── model.safetensors
    │   └── other files...
    ├── LlamaGen-T2I-2
    │   ├── config.json
    │   ├── generation_config.json
    │   ├── model.safetensors
    │   └── other files...
    ├── vq_ds16_t2i.pt
    └── t5
        └── flan-t5-xl
            ├── config.json
            ├── generation_config.json
            ├── model-00001-of-00002.safetensors
            └── other files...

      (Optional) Trained drafter To use trained drafter, you need to download llamagen_drafter and/or llamagen2_drafter and save it under trained_drafters directory.

      ckpts
└── llamagen
    └── trained_drafters
        ├── llamagen_drafter
        |   ├── config.json
        |   ├── generation_config.json
        |   ├── pytorch_model.bin
        |   └── other files...
        └── llamagen2_drafter
            ├── config.json
            ├── generation_config.json
            ├── pytorch_model.bin
            └── other files...

    3. Anole For Anole, download Anole-7b-v0.1-hf, which is a huggingface style converted model from Anole.

      In addition, you should download the original VQ-VAE weights provided by Meta and put them to the following directory:

      ckpts
└── anole
    ├── Anole-7b-v0.1-hf
    |   ├── config.json
    |   ├── generation_config.json
    |   ├── model-00001-of-00003.safetensors
    |   └── other files...
    └── chameleon
        └── tokenizer
            ├── text_tokenizer.json
            ├── vqgan.yaml
            └── vqgan.ckpt

      (Optional) Trained drafter To use trained drafter, you need to download anole_drafter and save it under trained_drafters directory.

      ckpts
└── anole
    └── trained_drafters
        └── anole_drafter
            ├── config.json
            ├── generation_config.json
            ├── pytorch_model.bin
            └── other files...

✨ Usage

All the functionalities can be done by either running main.py or directly running entrypoints/{function}.py. Currently, "llamagen" (LlamaGen-Stage I), "llamagen2" (LlamaGen-Stage II), "anole", and "lumina_mgpt" are supported as --model.

🚧 Lumina-mGPT is still under construction, so some functions may not work properly yet. You can follow the procedures here, but you may encounter a few exceptions.

  1. Generate Images

    python main.py generate_images --model <model_name> --model_type <model_type; e.g., base, vllm, eagle> --model_path <model_path> --drafter_path <drafter_path> --output_dir <output_dir> ...

    or

    python -m entrypoints.generate_images --model <model_name> --model_type <model_type; e.g., base, vllm, eagle> --model_path <model_path> --drafter_path <drafter_path> --output_dir <output_dir> ...

    💡How to use LANTERN and LANTERN++ for image generation

    • For LANTERN, set --model_type eagle, turn on --lantern option and set --lantern_k and --lantern_delta options.
    • For LANTERN++, use --static_tree option and use --lantern_delta to set $\lambda$ value.

  2. Generate Training Data for Drafter

    python main.py generate_train_data --model <model_name> --data_path <path_to_image_tokens> --output_dir <output_dir> --num_samples <num_samples>

    or

    python -m entrypoints.generate_train_data --model <model_name> --data_path <path_to_image_tokens> --output_dir <output_dir> --num_samples <num_samples>

    For LlamaGen and Anole, you have to extract code and T5 embedding(only for LlamaGen) for training data.

    • Locate image and caption files in given format and execute following command before run generate_train_data:

    Data Format:

    • image_folder
      • {file_1}.jpg
      • {file_1}.txt
      • {file_2}.jpg
      • {file_2}.txt ...
    python main.py extract_code --model <model_type> --data_path <path_to_image_and_caption> --output_dir <output_dir> --num_samples <num_samples>

    or

    python -m entrypoints.extract_code --model <model_type> --data_path <path_to_image_and_caption> --output_dir <output_dir> --num_samples <num_samples>

  3. Train Drafter Model

     python main.py train_drafter --model <model_type> --base_path <base_model_path> --config_path <path_to_config.json> --data_dir <data_dir> --save_dir <save_dir> --lr <lr> --bs <bs> --gradient_accumlation_steps <gradient_accumulation_steps> ...

    or

    python -m entrypoints.train_drafter.main --model <model_type> --base_path <base_model_path> --config_path <path_to_config.json> --data_dir <data_dir> --save_dir <save_dir> --lr <lr> --bs <bs> --gradient_accumlation_steps <gradient_accumulation_steps> ...

    For multi GPU training with accelerate, you can use

     accelerate launch -m entrypoints.train_drafter.main --model <model_type> --base_path <base_model_path> --config_path <path_to_config.json> --data_dir <data_dir> --save_dir <save_dir> --lr <lr> --bs <bs> --gradient_accumlation_steps <gradient_accumulation_steps> ...

  4. Generate VQ Distances

    python main.py generate_codebook --model <model_name> --save_path <save_path>

    or

    python -m entrypoints.generate_codebook --model <model_name> --save_path <save_path>

  5. Evaluate Generated Images We support FID, CLIP score, Precision/Recall and HPSv2 for image evaluation.

    python main.py eval_fid_clip --fake_dir <path_to_generated_image> --ref_dir <path_to_reference_image> --caption_path <path_to_prompt> --how_many <number_of_images_for_evaluation> ...
    python main.py eval_prec_recall --fake_dir <path_to_generated_image> --ref_dir <path_to_reference_image> ...
    python main.py eval_hpsv2 --image_path <path_to_generated_image> --prompt_path <path_to_prompt>

    or

    python -m entrypoints.eval_fid_clip --fake_dir <path_to_generated_image> --ref_dir <path_to_reference_image> --caption_path <path_to_prompt> --how_many <number_of_images_for_evaluation> ...
    python -m entrypoints.eval_prec_recall --fake_dir <path_to_generated_image> --ref_dir <path_to_reference_image> ...
    python -m entrypoints.eval_hpsv2 --image_path <path_to_generated_image> --prompt_path <path_to_prompt>

⚠️ CAUTIONS

  1. config.json should be in the ckpts/{model_name}/trained_models/{drafter_path} Since the Model in cnets_{model_name}.py is initialized according to the config.json in the drafter_path, you need to place config.json for drafter correctly. Note that the config.json should be same as the base model's config.json other than num_hidden_layers.

⚖️ License

This project is distributed under the Chameleon License by Meta Platforms, Inc. For more information, please see the LICENSE file in the repository.

🙏 Acknowledgement

This repository is built with extensive reference to FoundationVision/LlamaGen, Alpha-VLLM/Lumina-mGPT and SafeAILab/EAGLE, leveraging many of their core components and approaches.

📄 Citation

@article{jang2024lantern,
  title={LANTERN: Accelerating Visual Autoregressive Models with Relaxed Speculative Decoding},
  author={Jang, Doohyuk and Park, Sihwan and Yang, June Yong and Jung, Yeonsung and Yun, Jihun and Kundu, Souvik and Kim, Sung-Yub and Yang, Eunho},
  journal={arXiv preprint arXiv:2410.03355},
  year={2024}
}
@article{park2025lanternenhancedrelaxedspeculative,
  title={LANTERN++: Enhanced Relaxed Speculative Decoding with Static Tree Drafting for Visual Auto-regressive Models}, 
  author={Sihwan Park and Doohyuk Jang and Sungyub Kim and Souvik Kundu and Eunho Yang},
  journal={arXiv preprint arXiv:2410.03355},
  year={2025}
}