DOCSEARCHbyAlgolia
Submit Docs

EfficientViT

repository·master·Indexed 25 days ago

https://github.com/mit-han-lab/efficientvit

A suite of efficient vision foundation models for high-resolution generation and perception tasks, including segmentation, classification, and diffusion-based image generation. The library includes the Deep Compression Autoencoder (DC-AE) for image encoding and decoding, compatible with the diffusers ecosystem, and provides tools for training and evaluating classification models and DC-AE-Diffusion models (USiT, UViT, DiT). It supports exporting models to ONNX and TFLite formats.

Tokens
14.5K
Snippets
36
Records
78
Agent score
85%

What's inside EfficientViT

  1. Overview of EfficientViT Applications

    master

    The EfficientViT repository contains several specialized vision foundation models:

    • Deep Compression Autoencoder (DC-AE): High-spatial compression autoencoders (up to 128x) designed to accelerate latent diffusion models.
    • EfficientViT-SAM: An accelerated Segment Anything Model (SAM) that replaces the heavy image encoder with EfficientViT, providing significant speedups (e.g., 48.9x on A100 via TensorRT) without accuracy loss.
    • EfficientViT-Classification: Lightweight image classification models using EfficientViT backbones.
    • EfficientViT-Segmentation: Efficient semantic segmentation models.
    • EfficientViT-GazeSAM: Gaze-prompted image segmentation models capable of real-time execution with TensorRT.

  3. Use Deep Compression Autoencoder (DC-AE) with Diffusers

    master

    You can use DC-AE models within the diffusers ecosystem by installing the diffusers library. The AutoencoderDC class allows you to load pretrained DC-AE models from Hugging Face to encode images into latents and decode them back.

    pip install -U diffusers
    from PIL import Image
import torch
import torchvision.transforms as transforms
from torchvision.utils import save_image
from diffusers import AutoencoderDC

device = torch.device("cuda")
dc_ae: AutoencoderDC = AutoencoderDC.from_pretrained(f"mit-han-lab/dc-ae-f64c128-in-1.0-diffusers", torch_dtype=torch.float32).to(device).eval()

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(0.5, 0.5),
])

image = Image.open("assets/fig/girl.png")
x = transform(image)[None].to(device)
latent = dc_ae.encode(x).latent
y = dc_ae.decode(latent).sample
save_image(y * 0.5 + 0.5, "demo_dc_ae.png")

  4. Evaluate DC-AE-Diffusion Models

    master

    To evaluate DC-AE-Diffusion models, you must first generate a reference for FID (Fréchet Inception Distance) computation, then run the evaluation script.

    1. Generate FID Reference

    Use generate_reference to create the necessary .npz file for FID calculation.

    2. Run Evaluation

    Use eval_dc_ae_diffusion_model. You can run evaluation with or without Classifier-Free Guidance (CFG) by setting the cfg_scale parameter.

    # 1. Generate reference for FID computation
torchrun --nnodes=1 --nproc_per_node=8 -m applications.dc_ae.generate_reference \
    dataset=imagenet imagenet.resolution=512 imagenet.image_mean=[0.,0.,0.] imagenet.image_std=[1.,1.,1.] split=train \
    fid.save_path=assets/data/fid/imagenet_512_train.npz

# 2. Run evaluation without cfg (cfg_scale=1.0)
torchrun --nnodes=1 --nproc_per_node=8 -m applications.dc_ae.eval_dc_ae_diffusion_model dataset=imagenet_512 model=dc-ae-f64c128-in-1.0-uvit-h-in-512px cfg_scale=1.0 run_dir=tmp

# 3. Run evaluation with cfg (e.g., cfg_scale=1.5)
torchrun --nnodes=1 --nproc_per_node=8 -m applications.dc_ae.eval_dc_ae_diffusion_model dataset=imagenet_512 model=dc-ae-f64c128-in-1.0-uvit-h-in-512px cfg_scale=1.5 run_dir=tmp

  5. Install EfficientViT

    master

    To set up the EfficientViT environment, create a new Conda environment with Python 3.10 and install the required dependencies using the provided requirements.txt file.

    conda create -n efficientvit python=3.10
conda activate efficientvit
pip install -U -r requirements.txt

  8. Generate images with DC-AE Diffusion Models

    master

    To perform text-to-image generation using DC-AE integrated diffusion models, use the DCAE_Diffusion_HF class from efficientvit.diffusion_model_zoo. The workflow involves generating latents in the latent space using the diffusion model and then decoding those latents using the attached autoencoder.

    # build DC-AE-Diffusion models
from efficientvit.diffusion_model_zoo import DCAE_Diffusion_HF

dc_ae_diffusion = DCAE_Diffusion_HF.from_pretrained(f"mit-han-lab/dc-ae-f64c128-in-1.0-uvit-h-in-512px-train2000k")

# denoising on the latent space
import torch
import numpy as np
from torchvision.utils import save_image

torch.set_grad_enabled(False)
device = torch.device("cuda")
dc_ae_diffusion = dc_ae_diffusion.to(device).eval()

seed = 0
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
eval_generator = torch.Generator(device=device)
eval_generator.manual_seed(seed)

prompts = torch.tensor(
    [279, 333, 979, 936, 933, 145, 497, 1, 248, 360, 793, 12, 387, 437, 938, 978], dtype=torch.int, device=device
)
num_samples = prompts.shape[0]
prompts_null = 1000 * torch.ones((num_samples,), dtype=torch.int, device=device)
latent_samples = dc_ae_diffusion.diffusion_model.generate(prompts, prompts_null, 6.0, eval_generator)
latent_samples = latent_samples / dc_ae_diffusion.scaling_factor

# decode
image_samples = dc_ae_diffusion.autoencoder.decode(latent_samples)
save_image(image_samples * 0.5 + 0.5, "demo_dc_ae_diffusion.png", nrow=int(np.sqrt(num_samples)))

  9. Generate optimized TensorRT engines (INT8)

    master

    To generate optimized TensorRT engines using INT8 precision for specific models, follow these steps. The image encoder remains FP32, the image decoder and depth/face models use FP16, and the gaze estimation and object detection models use INT8.

    1. Navigate to the models directory and create the cache directory: cd models; mkdir -p tensorrt/int8/caches
    2. Download INT8 calibration caches from Hugging Face and save them to tensorrt/int8/caches.
    3. If filenames contain the gazesam_int8_calib_caches_ prefix, remove it by running: rename 's/^gazesam_int8_calib_caches_//' gazesam_int8_calib_caches_*.cache (while inside tensorrt/int8/caches).
    4. Execute the optimized creation script: bash create_optimized_engines.sh
    cd models; mkdir -p tensorrt/int8/caches
# Download caches to models/tensorrt/int8/caches/
# Optional: rename files to remove prefix
rename 's/^gazesam_int8_calib_caches_//' gazesam_int8_calib_caches_*.cache
bash create_optimized_engines.sh

  10. Train EfficientViT L Series models

    master

    Train the L series models using torchrun. You need to provide the configuration YAML, enable mixed precision with --amp bf16, and specify the data directory and output path.

    # Example for EfficientViT-L1
torchrun --nnodes 1 --nproc_per_node=8 \
python applications/efficientvit_cls/train_efficientvit_cls_model.py applications/efficientvit_cls/configs/imagenet/efficientvit_l1.yaml --amp bf16 \
    --data_provider.data_dir ~/dataset/imagenet \
    --path .exp/efficientvit_cls/imagenet/efficientvit_l1_r224/

  11. Set up EfficientViT-GazeSAM with ONNX

    master

    To use the ONNX runtime:

    1. Install ONNX Runtime GPU:

      python -m pip install onnxruntime-gpu

      Note: If you encounter issues, try uninstalling both onnxruntime and onnxruntime-gpu, then reinstalling only onnxruntime-gpu.

    2. Download Model Components: Download the ONNX model components from Hugging Face and save them to the models/onnx directory (ensure the onnx subfolder exists).

    python -m pip install onnxruntime-gpu

  12. Extract latents using Flux VAE

    master

    Extract latents from an image dataset using the applications.dc_ae.dc_ae_generate_latent module. This requires torchrun for distributed execution.

    torchrun --nnodes 1 --nproc_per_node=8 -m applications.dc_ae.dc_ae_generate_latent resolution=512 \
    image_root_path=~/dataset/imagenet/train batch_size=64 \
    model_name=flux-vae \
    latent_root_path=assets/data/latent/flux_vae/imagenet_512