# VLM Inference Guide

Guide for running inference and chatbots with Vision-Language Models (AR and MDM).

## Inference Scripts

| Script | Model Type | Description |
|--------|------------|-------------|
| `chatbot.py` | AR | Standard autoregressive chatbot |
| `chatbot_mdm.py` | MDM | Masked diffusion model chatbot |
| `chatbot_mdm_v2.py` | MDM | MDM chatbot with KV cache |
| `run_qwen_vl_chatbot_text_only.py` | AR | Text-only inference |
| `run_qwen_vl_chatbot_text_only_mdm.py` | MDM | Text-only MDM inference |

---

## MDM Chatbot

### Basic Usage

```bash
python chengyue_script/dvlm/inference/chatbot_mdm.py \
  --model-name output_models/j63_mdm/checkpoint-2000 \
  --processor-name Qwen/Qwen2.5-VL-3B-Instruct \
  --demo
```

### Command Line Arguments

| Argument | Description | Default |
|----------|-------------|---------|
| `--model-name` | Model checkpoint path | Required |
| `--processor-name` | Processor/tokenizer path | `Qwen/Qwen2.5-VL-3B-Instruct` |
| `--image` | Image URL or path | None |
| `--prompt` | Text prompt | None |
| `--demo` | Run demo samples | False |
| `--block-size` | Block size for MDM sampling | 8 |
| `--max-tokens` | Maximum tokens to generate | 512 |
| `--threshold` | Confidence threshold for unmasking | 0.9 |
| `--device` | Device (cuda/cpu) | cuda |

### Custom Input

```bash
python chengyue_script/dvlm/inference/chatbot_mdm.py \
  --model-name output_models/j63_mdm/checkpoint-2000 \
  --image "path/to/image.jpg" \
  --prompt "Describe this image in detail."
```

### Demo Mode

Demo mode runs predefined samples from `demo_images/`:

```bash
python chengyue_script/dvlm/inference/chatbot_mdm.py \
  --model-name output_models/j63_mdm/checkpoint-2000 \
  --demo
```

Demo samples include:
- Beach dog image with description prompt
- MathVista physics problem
- MathVerse geometry problem
- MathVision number puzzle

---

## MDM Sampling Algorithm

The MDM sampling process uses block-wise denoising:

```
Input: prompt tokens + [MASK] * block_size
  │
  ▼
┌─────────────────────────────────────┐
│ For each block:                     │
│   1. Initialize block with [MASK]   │
│   2. While any [MASK] in block:     │
│      - Forward pass to get logits   │
│      - Compute confidence scores    │
│      - Unmask high-confidence tokens│
│      - If none > threshold:         │
│        unmask highest confidence    │
│   3. Move to next block             │
└─────────────────────────────────────┘
  │
  ▼
Output: generated tokens
```

### Key Parameters

- **`block_size`**: Number of tokens processed in parallel (default: 8)
- **`threshold`**: Confidence threshold for unmasking (default: 0.9)
- **`max_tokens`**: Maximum generation length (default: 512)

---

## AR Chatbot

### Basic Usage

```bash
python chengyue_script/dvlm/inference/chatbot.py \
  --model-name output_models/ar_baseline/checkpoint-2000 \
  --image "path/to/image.jpg" \
  --prompt "What is in this image?"
```

---

## Programmatic Usage

### Using QwenVLMDMModel Class

```python
from chengyue_script.dvlm.models.qwen_vl_mdm_model import QwenVLMDMModel

# Initialize model
model = QwenVLMDMModel(
    model_name="output_models/j63_mdm/checkpoint-2000",
    processor_name="Qwen/Qwen2.5-VL-3B-Instruct",
    block_size=8,
    max_tokens=128,
)

# Generate response
response = model.generate(
    prompt="Describe this image.",
    image="path/to/image.jpg",
)
print(response)

# Batch generation
responses = model.generate_batch(
    prompts=["Describe image 1.", "Describe image 2."],
    images=["image1.jpg", "image2.jpg"],
)
```

### Direct Model Usage

```python
import torch
from transformers import AutoModel, AutoProcessor, AutoTokenizer
from chengyue_script.dvlm.models.generation_functions import Fast_dLLM_QwenForCausalLM
import types

# Load model
model = AutoModel.from_pretrained(
    "output_models/j63_mdm/checkpoint-2000",
    torch_dtype=torch.bfloat16,
    device_map="cuda",
    trust_remote_code=True,
)

# Bind MDM sampling method
model.mdm_sample = types.MethodType(Fast_dLLM_QwenForCausalLM.mdm_sample, model)

# Load processor
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("output_models/j63_mdm/checkpoint-2000")
processor.tokenizer = tokenizer

# Prepare input
messages = [{"role": "user", "content": [
    {"type": "image", "image": "path/to/image.jpg"},
    {"type": "text", "text": "Describe this image."},
]}]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=[text], images=[image], return_tensors="pt").to("cuda")

# Generate
mask_id = tokenizer.encode("|<MASK>|")[0]
output_ids = model.mdm_sample(
    input_ids=inputs.input_ids,
    tokenizer=tokenizer,
    block_size=8,
    max_tokens=128,
    mask_id=mask_id,
    pixel_values=inputs.pixel_values,
    image_grid_thw=inputs.image_grid_thw,
)

# Decode
response = processor.batch_decode(output_ids[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0]
```

---

## KV Cache Optimization

`chatbot_mdm_v2.py` uses KV cache for faster generation:

```bash
python chengyue_script/dvlm/inference/chatbot_mdm_v2.py \
  --model-name output_models/j63_mdm/checkpoint-2000 \
  --demo
```

The KV cache version:
1. Caches prompt tokens in first forward pass
2. Only processes current block tokens in subsequent passes
3. Significantly faster for long prompts

---

## Demo Images

Demo images are located in `dvlm/demo_images/`:

| Image | Description |
|-------|-------------|
| `beach_dog.jpeg` | Dog on beach (general VQA) |
| `mathvista_physics.png` | Physics spring problem |
| `mathverse.jpg` | Geometry triangle problem |
| `mathvision.jpg` | Number sequence puzzle |

---

## Troubleshooting

### CUDA Out of Memory

1. Reduce `--max-tokens`
2. Use smaller images
3. Set `--device cpu` for CPU inference

### Model Loading Errors

1. Ensure checkpoint path is correct
2. Verify processor is from base model: `Qwen/Qwen2.5-VL-3B-Instruct`
3. Check `trust_remote_code=True` is set

### Slow Generation

1. Use KV cache version (`chatbot_mdm_v2.py`)
2. Reduce `--max-tokens`
3. Increase `--block-size` (trades quality for speed)

