# How Image Conditioning Works with the `ImageFeatureExtractor` in TRELLIS 2

> Learn how TRELLIS 2 uses DINO-based feature extractors for image conditioning. Discover how input images become patch-level embeddings and inject into generative models via a lazy-initialization mixin.

- Repository: [Microsoft/TRELLIS.2](https://github.com/microsoft/TRELLIS.2)
- Tags: how-to-guide
- Published: 2026-08-04

---

**TRELLIS 2 enables image-conditioned generative 3D models by converting input images into patch-level embeddings through DINO-based feature extractors, then injecting these conditioning vectors into the diffusion or flow-matching backbone via a lazy-initialization mixin pattern.**

Image conditioning is central to TRELLIS 2's ability to generate 3D assets guided by visual input. The architecture separates feature extraction from model training through clean abstractions: specialized extractors handle vision transformer backbones, while mixins transparently add conditioning capability to any trainer. This design follows the repository's modular philosophy in `microsoft/TRELLIS.2`.

## How the `ImageFeatureExtractor` Transforms Images into Conditioning Vectors

TRELLIS 2 provides two concrete implementations in [`trellis2/modules/image_feature_extractor.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/image_feature_extractor.py): `DinoV2FeatureExtractor` and `DinoV3FeatureExtractor`. Both follow the same six-step pipeline to produce embeddings suitable for generative conditioning.

### Loading the DINO Backbone

Each extractor initializes its vision transformer through the standard model hub:

- **DINO v2**: Loaded via `torch.hub.load('facebookresearch/dinov2', model_name, pretrained=True)` at lines 16-18
- **DINO v3**: Loaded via `DINOv3ViTModel.from_pretrained` for the newer architecture

This lazy-download approach ensures pretrained weights are fetched only when needed.

### Image Preprocessing and Normalization

Raw PIL images undergo standardized transformation before reaching the backbone:

```python
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

```

This ImageNet normalization matches the statistics used during DINO pretraining. The full preprocessing pipeline (lines 46-51 for v2, 108-113 for v3) resizes images, converts to RGB, stacks into batched tensors, and moves to GPU.

### Feature Extraction and Layer Normalization

The forward pass yields raw patch tokens that require stabilization:

```python
features = self.model(image, is_training=True)['x_prenorm']
patchtokens = F.layer_norm(features, features.shape[-1:])  # lines 55-57

```

The `is_training=True` flag preserves gradient flow through the backbone when needed. Layer normalization operates on the final dimension independently for each patch token, preventing distribution shift across different input images.

### Output Format for Downstream Consumption

Both extractors return a tensor of shape **(B, N, D)** where:
- **B** = batch size
- **N** = number of image patches (e.g., 256 for a 224×224 image with 14×14 patches)
- **D** = feature dimension (384 for DINO-S, 768 for DINO-B, 1024 for DINO-L)

This format directly feeds into cross-attention or concatenation-based conditioning mechanisms in the generative backbone.

## The `ImageConditionedMixin`: Seamless Integration with Trainers

The `ImageConditionedMixin` in [`trellis2/trainers/flow_matching/mixins/image_conditioned.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/trainers/flow_matching/mixins/image_conditioned.py) eliminates boilerplate when adding image conditioning to any trainer. It handles three critical concerns: lazy initialization, distributed-safe setup, and API compatibility.

### Lazy Initialization with Configuration-Driven Setup

The mixin receives a configuration dictionary rather than an instantiated model:

```python
self.image_cond_model_config = {'name': 'DinoV2FeatureExtractor', 
                                 'args': {'model_name': 'dinov2_vits14'}}

```

The actual extractor is constructed only on first use inside `_init_image_cond_model` (lines 35-41), wrapped in `dist_utils.local_master_first()` to prevent race conditions when multiple GPU processes attempt simultaneous downloads:

```python
with dist_utils.local_master_first():
    self.image_cond_model = globals()[self.image_cond_model_config['name']](
        **self.image_cond_model_config.get('args', {})
    )

```

### Encoding Images to Conditioning Tensors

The `encode_image` method (lines 48-51) provides a simple interface:

```python
features = self.image_cond_model(image)  # PIL Image → (1, N, D) tensor

```

For multi-view conditioning, `MultiImageConditionedMixin` provides `encode_images` with analogous behavior.

### Transparent Conditioning Injection

The mixin overrides three key trainer methods to maintain API compatibility:

1. **`get_cond`** — Encodes the image, creates zero-filled negative conditioning, and forwards to base class (lines 54-60)
2. **`get_inference_cond`** — Same pattern for sampling/inference
3. **`vis_cond`** — Visualization hook for debugging conditioning inputs

This design means trainers inherit conditioning capability without modifying their core training loops. The negative conditioning (`torch.zeros_like(cond)`) provides a clean fallback for classifier-free guidance implementations.

## End-to-End Usage Example

Here's a complete workflow showing how image conditioning integrates into training and inference:

```python
from trellis2.trainers.flow_matching import SparseFlowMatching
from PIL import Image
import torch

# ----------------------------------------------------------------------

# Configure image conditioning

# ----------------------------------------------------------------------

image_cond_cfg = {
    "name": "DinoV2FeatureExtractor",
    "args": {"model_name": "dinov2_vits14"}  # Small model, 384-dim features

}

# ----------------------------------------------------------------------

# Instantiate trainer with conditioning support

# ----------------------------------------------------------------------

trainer = SparseFlowMatching(
    model_config={...},           # Your flow-matching architecture

    image_cond_model=image_cond_cfg,  # Enables ImageConditionedMixin

)

# ----------------------------------------------------------------------

# Direct encoding (useful for debugging or external pipelines)

# ----------------------------------------------------------------------

pil_img = Image.open("reference.jpg").convert("RGB")
cond_tensor = trainer.encode_image(pil_img)
print(cond_tensor.shape)  # torch.Size([1, 256, 384])

# ----------------------------------------------------------------------

# Training with automatic conditioning injection

# ----------------------------------------------------------------------

# The mixin automatically encodes `cond` and handles neg_cond internally

loss = trainer.train_step(batch_data, cond=pil_img)

# ----------------------------------------------------------------------

# Sampling with image guidance

# ----------------------------------------------------------------------

samples = trainer.sample(num_steps=100, cond=pil_img)

```

## Multi-Image Conditioning for View-Consistent Generation

`MultiImageConditionedMixin` extends this pattern to handle multiple input images—critical for 3D generation from sparse views:

```python

# Configuration for multi-image conditioning

multi_cfg = {
    "name": "DinoV2FeatureExtractor",
    "args": {"model_name": "dinov2_vitb14"}  # Base model for richer features

}

# Encode multiple views

views = [Image.open(f"view_{i}.jpg") for i in range(4)]
cond_tensor = trainer.encode_images(views)  # Shape: (4, N, D) or fused representation

```

The mixin handles view fusion strategies internally, either stacking features or applying learned aggregation depending on the specific trainer implementation.

## Architecture Benefits of This Design

- **Modularity**: Swapping DINO v2 for v3 requires only a configuration change, no code modification
- **Memory efficiency**: Lazy initialization prevents GPU memory allocation until the first forward pass
- **Distributed safety**: `local_master_first()` ensures only one process downloads weights in multi-GPU setups
- **Backward compatibility**: Non-conditioned trainers remain unchanged; conditioning is opt-in via inheritance

## Key Source Files

| File | Purpose | Key Components |
|------|---------|---------------|
| [`trellis2/modules/image_feature_extractor.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/modules/image_feature_extractor.py) | DINO-based feature extraction | `DinoV2FeatureExtractor`, `DinoV3FeatureExtractor` |
| [`trellis2/trainers/flow_matching/mixins/image_conditioned.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/trainers/flow_matching/mixins/image_conditioned.py) | Conditioning integration | `ImageConditionedMixin`, `MultiImageConditionedMixin` |
| [`trellis2/utils/dist_utils.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/utils/dist_utils.py) | Distributed training utilities | `local_master_first()` |
| [`trellis2/trainers/flow_matching/sparse_flow_matching.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/trainers/flow_matching/sparse_flow_matching.py) | Example conditioned trainer | Inherits mixin functionality |

## Summary

- **Image conditioning** in TRELLIS 2 converts input images to patch-level embeddings using **DINO v2 or v3 vision transformers**
- The **`ImageFeatureExtractor`** implementations in [`image_feature_extractor.py`](https://github.com/microsoft/TRELLIS.2/blob/main/image_feature_extractor.py) handle preprocessing, normalization, and layer-normalized feature extraction
- **`ImageConditionedMixin`** provides lazy initialization, distributed-safe setup, and transparent API integration via configuration dictionaries
- Conditioning tensors follow shape **(B, N, D)** and are automatically paired with zero negative conditioning for flexible guidance strategies
- **Multi-image support** extends the same pattern through `MultiImageConditionedMixin` for multi-view 3D generation scenarios

## Frequently Asked Questions

### What DINO model sizes are supported by the `ImageFeatureExtractor`?

TRELLIS 2 supports all official DINO v2 and v3 variants including `dinov2_vits14` (small, 384-dim), `dinov2_vitb14` (base, 768-dim), `dinov2_vitl14` (large, 1024-dim), and `dinov2_vitg14` (giant, 1536-dim). The model name is passed directly to `torch.hub.load` or `from_pretrained`, so any future DINO release with compatible output structure will work without code changes—only configuration modification is required.

### Why does the mixin use lazy initialization instead of creating the extractor in `__init__`?

Lazy initialization prevents unnecessary GPU memory allocation and weight downloads when the trainer is instantiated for operations that don't require image conditioning, such as loading checkpoints for inference on non-conditioned tasks or running unit tests. It also allows `local_master_first()` to coordinate initialization in distributed settings, avoiding race conditions where multiple processes simultaneously attempt to download the same backbone weights.

### How can I extract features from multiple images for a single conditioning vector?

Use `MultiImageConditionedMixin` instead of `ImageConditionedMixin` and call `encode_images(view_list)` rather than `encode_image(single_image)`. The multi-image variant handles batching and fusion internally, typically returning either stacked features or aggregated representations depending on your trainer's specific implementation of the conditioning pathway.

### Is the DINO backbone frozen during training or fine-tuned?

The analysis shows `is_training=True` being passed to the DINO forward pass, which preserves gradients through the backbone. However, typical TRELLIS 2 configurations freeze these pretrained weights and use them as a fixed feature extractor to reduce compute and prevent overfitting. Check your specific trainer configuration for `requires_grad` settings on the `image_cond_model` parameters if you need fine-tuning behavior.