# VLM Model Utilities

This directory contains model utilities and generation functions for Vision-Language Models.

## Files

| File | Description |
|------|-------------|
| `generation_functions.py` | MDM sampling functions for text generation |
| `qwen_vl_mdm_model.py` | High-level wrapper for MDM inference |

---

## Generation Functions

### Fast_dLLM_QwenForCausalLM

Core class providing MDM (Masked Diffusion Model) sampling methods:

#### `mdm_sample`

Standard MDM sampling without KV cache:

```python
def mdm_sample(
    self,
    input_ids,           # Input token IDs
    tokenizer,           # Tokenizer for decoding
    block_size,          # Block size for masking (default: 8)
    max_tokens,          # Maximum tokens to generate
    pixel_values=None,   # Image pixel values (optional)
    image_grid_thw=None, # Image grid dimensions (optional)
    mask_id=151665,      # Mask token ID
    threshold=0.9,       # Confidence threshold for unmasking
    output_file='decoded_output.txt',  # Debug output file
    stop_token=151645,   # Stop token ID
):
```

#### `mdm_sample_cache`

MDM sampling with KV cache for faster generation:

```python
def mdm_sample_cache(
    self,
    input_ids,
    tokenizer,
    block_size,
    max_tokens, 
    pixel_values=None,
    image_grid_thw=None,
    mask_id=151665,
    threshold=0.9,
    stop_token=151645,
):
```

---

## MDM Sampling Algorithm

### Block-wise Denoising

MDM generates tokens in blocks rather than one-by-one:

```
Block 1: [MASK][MASK][MASK][MASK][MASK][MASK][MASK][MASK]
    ↓ Forward pass
Block 1: [token][MASK][token][MASK][MASK][token][MASK][MASK]  (high-confidence unmasked)
    ↓ Forward pass
Block 1: [token][token][token][MASK][token][token][MASK][MASK]
    ↓ Forward pass
Block 1: [token][token][token][token][token][token][token][token]  (all unmasked)
    ↓
Block 2: [MASK][MASK][MASK][MASK][MASK][MASK][MASK][MASK]
    ... (repeat)
```

### Unmasking Strategy

For each forward pass within a block:

1. **Compute confidence**: `softmax(logits)` for each position
2. **High-confidence unmasking**: Unmask all tokens with confidence > `threshold`
3. **Fallback**: If no tokens exceed threshold, unmask the highest-confidence token
4. **Repeat** until all tokens in block are unmasked

---

## QwenVLMDMModel

High-level wrapper class for easy inference:

```python
from models.qwen_vl_mdm_model import QwenVLMDMModel

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

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

# Batch generation
responses = model.generate_batch(
    prompts=["Q1", "Q2"],
    images=["img1.jpg", "img2.jpg"],
)
```

---

## Integration with Training

During training, the model uses a specialized attention mask for block diffusion:

### Block Diffusion Attention Mask

The attention mask combines three components:

1. **Block Diagonal Mask (M_BD)**: Self-attention within noised blocks
2. **Offset Block Causal Mask (M_OBC)**: Cross-attention for conditional context
3. **Block Causal Mask (M_BC)**: Attention to update x0

```python
def block_diff_mask(b, h, q_idx, kv_idx, block_size=None, n=None):
    # Block indices
    block_q = q_idx // block_size
    block_kv = kv_idx // block_size
    
    # Combine masks
    block_diagonal = (block_q == block_kv) & (x0_flag_q == x0_flag_kv)
    offset_block_causal = (block_q > block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 0)
    block_causal = (block_q >= block_kv) & (x0_flag_kv == 1) & (x0_flag_q == 1)
    
    return block_diagonal | offset_block_causal | block_causal
```

---

## Usage Example

### Binding to Model

```python
import types
from transformers import AutoModel
from generation_functions import Fast_dLLM_QwenForCausalLM

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

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

# Or with KV cache
model.mdm_sample = types.MethodType(Fast_dLLM_QwenForCausalLM.mdm_sample_cache, model)
```

### Setup Function

```python
from generation_functions import setup_model_with_custom_generation

model = AutoModel.from_pretrained(...)
model = setup_model_with_custom_generation(model)
# Now model.mdm_sample is available
```

---

## Related Files

| File | Location | Purpose |
|------|----------|---------|
| `modeling_dfm.py` | `chengyue_script/` | Full DFM model implementation |
| `local_modeling_dfm.py` | `src/models/Qwen2_5VL_MDM/` | DFM model for HF loading |
| `modeling_qwen2vl_mdm.py` | `src/models/Qwen2_5VL_MDM/` | MDM-Qwen2VL wrapper |

