179 lines
8.2 KiB
Markdown
179 lines
8.2 KiB
Markdown
|
# BAGEL-7B-MoT
|
|||
|
|
|
|||
|
|
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/bagel>.
|
|||
|
|
|
|||
|
|
## Set up
|
|||
|
|
|
|||
|
|
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
|||
|
|
|
|||
|
|
## Run examples
|
|||
|
|
|
|||
|
|
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
|
|||
|
|
|
|||
|
|
Get into the bagel folder
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
cd examples/offline_inference/bagel
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
### Modality Control
|
|||
|
|
|
|||
|
|
BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
|
|||
|
|
|
|||
|
|
#### Text to Image (text2img)
|
|||
|
|
|
|||
|
|
- **Pipeline**: Text → Thinker → DiT → VAE Decode → Image
|
|||
|
|
- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
|
|||
|
|
- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
|
|||
|
|
|
|||
|
|
Generate images from text prompts:
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
|||
|
|
--modality text2img \
|
|||
|
|
--prompts "A cute cat"
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
#### Image to Image (img2img)
|
|||
|
|
|
|||
|
|
- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
|
|||
|
|
- **Stages Used**: Stage 1 (DiT) only
|
|||
|
|
- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
|
|||
|
|
|
|||
|
|
Transform images based on text prompts:
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
|||
|
|
--modality img2img \
|
|||
|
|
--image-path /path/to/image.jpg \
|
|||
|
|
--prompts "Let the woman wear a blue dress"
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
#### Image to Text (img2text)
|
|||
|
|
|
|||
|
|
- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
|
|||
|
|
- **Stages Used**: Stage 0 (Thinker) only
|
|||
|
|
- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
|
|||
|
|
|
|||
|
|
Generate text descriptions from images:
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
|||
|
|
--modality img2text \
|
|||
|
|
--image-path /path/to/image.jpg \
|
|||
|
|
--prompts "Describe this image in detail"
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
#### Text to Text (text2text)
|
|||
|
|
|
|||
|
|
- **Pipeline**: Text → Thinker → Text Output
|
|||
|
|
- **Stages Used**: Stage 0 (Thinker) only
|
|||
|
|
- **Special**: No visual components involved, operates as pure language model
|
|||
|
|
|
|||
|
|
Pure text generation:
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
|||
|
|
--modality text2text \
|
|||
|
|
--prompts "What is the capital of France?"
|
|||
|
|
|
|||
|
|
# You can load prompts from a text file (one prompt per line):
|
|||
|
|
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
|||
|
|
--modality text2text \
|
|||
|
|
--txt-prompts /path/to/prompts.txt
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
### Inference Steps
|
|||
|
|
|
|||
|
|
Control the number of inference steps for image generation:
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
# You can adjust steps to 100 to improve image quality
|
|||
|
|
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
|||
|
|
--modality text2img \
|
|||
|
|
--steps 50 \
|
|||
|
|
--prompts "A cute cat"
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
### Key arguments
|
|||
|
|
|
|||
|
|
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
|
|||
|
|
|
|||
|
|
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
|
|||
|
|
|
|||
|
|
#### 📌 Command Line Arguments (end2end.py)
|
|||
|
|
|
|||
|
|
| Argument | Type | Default | Description |
|
|||
|
|
| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
|
|||
|
|
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name |
|
|||
|
|
| `--modality` | choice | `text2img` | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
|
|||
|
|
| `--prompts` | list | `None` | Input text prompts directly |
|
|||
|
|
| `--txt-prompts` | string | `None` | Path to txt file with one prompt per line |
|
|||
|
|
| `--image-path` | string | `None` | Input image path (for `img2img`/`img2text`) |
|
|||
|
|
| `--steps` | int | `50` | Number of inference steps |
|
|||
|
|
| `--stage-configs-path` | string | `None` | Custom stage config file path |
|
|||
|
|
| `--worker-backend` | choice | `process` | Worker backend: `process` or `ray` |
|
|||
|
|
| `--ray-address` | string | `None` | Ray cluster address |
|
|||
|
|
| `--enable-stats` | flag | `False` | Enable statistics logging |
|
|||
|
|
| `--init-sleep-seconds` | int | `20` | Initialization sleep time |
|
|||
|
|
| `--batch-timeout` | int | `5` | Batch timeout |
|
|||
|
|
| `--init-timeout` | int | `300` | Initialization timeout |
|
|||
|
|
|
|||
|
|
------
|
|||
|
|
|
|||
|
|
#### ⚙️ Stage Configuration Parameters (bagel.yaml)
|
|||
|
|
|
|||
|
|
**Stage 0 - Thinker (LLM Stage)**
|
|||
|
|
|
|||
|
|
| Parameter | Value | Description |
|
|||
|
|
| :------------------------------- | :------------------------------ | :----------------------- |
|
|||
|
|
| `stage_type` | `llm` | Stage type |
|
|||
|
|
| `devices` | `"0"` | GPU device ID |
|
|||
|
|
| `max_batch_size` | `1` | Maximum batch size |
|
|||
|
|
| `model_stage` | `thinker` | Model stage identifier |
|
|||
|
|
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
|
|||
|
|
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
|
|||
|
|
| `tensor_parallel_size` | `1` | Tensor parallel size |
|
|||
|
|
| `max_num_batched_tokens` | `32768` | Maximum batched tokens |
|
|||
|
|
| `omni_kv_config.need_send_cache` | `true` | Whether to send KV cache |
|
|||
|
|
|
|||
|
|
------
|
|||
|
|
|
|||
|
|
**Stage 1 - DiT (Diffusion Stage)**
|
|||
|
|
|
|||
|
|
| Parameter | Value | Description |
|
|||
|
|
| :------------------------------- | :---------- | :-------------------------- |
|
|||
|
|
| `stage_type` | `diffusion` | Stage type |
|
|||
|
|
| `devices` | `"0"` | GPU device ID |
|
|||
|
|
| `max_batch_size` | `1` | Maximum batch size |
|
|||
|
|
| `model_stage` | `dit` | Model stage identifier |
|
|||
|
|
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
|
|||
|
|
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
|
|||
|
|
| `engine_input_source` | `[0]` | Input source from Stage 0 |
|
|||
|
|
|
|||
|
|
------
|
|||
|
|
|
|||
|
|
#### 🔗 Runtime Configuration
|
|||
|
|
|
|||
|
|
| Parameter | Value | Description |
|
|||
|
|
| :-------------------- | :------ | :------------------------------- |
|
|||
|
|
| `window_size` | `-1` | Window size (-1 means unlimited) |
|
|||
|
|
| `max_inflight` | `1` | Maximum inflight requests |
|
|||
|
|
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
|
|||
|
|
|
|||
|
|
## FAQ
|
|||
|
|
|
|||
|
|
- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
sudo apt update
|
|||
|
|
sudo apt install ffmpeg
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
|
|||
|
|
|
|||
|
|
| Stage | VRAM |
|
|||
|
|
| :------------------ | :--------------------------- |
|
|||
|
|
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
|
|||
|
|
| Stage-1 (DiT) | **26.50 GiB** |
|
|||
|
|
| Total | **~42 GiB + KV Cache** |
|