How to Implement Batch Generation with Multiple Input Images in TRELLIS.2

TRELLIS.2 supports batch processing by passing a list of PIL images to the pipeline's run method, which leverages the built-in image_cond_model and sampler infrastructure to generate multiple 3D assets in a single forward pass.

While microsoft/TRELLIS.2 generates 3D geometry and texture from a single image by default, the underlying architecture in trellis2/pipelines/trellis2_image_to_3d.py already supports batched conditioning and sampling. By utilizing the DinoV2FeatureExtractor (or DinoV3FeatureExtractor) and setting the appropriate batch_size in sampler parameters, you can process multiple input images efficiently without external loops.

Prerequisites and Pipeline Initialization

Load the pretrained pipeline using the standard entry point. The Trellis2ImageTo3DPipeline class handles device placement and model initialization automatically.

from trellis2.pipelines.trellis2_image_to_3d import Trellis2ImageTo3DPipeline

pipeline = Trellis2ImageTo3DPipeline.from_pretrained(
    path="path/to/pretrained/model"
)

Implementing Batch Generation

Method 1: Native Batch Processing (Recommended)

Pass a list of PIL.Image objects directly to the image argument. The pipeline's image_cond_model accepts batched inputs and returns a feature tensor of shape (B, N, D), where B is the batch size.

from PIL import Image

# Load multiple input images

images = [
    Image.open("data/chair.png"),
    Image.open("data/table.png"),
    Image.open("data/lamp.png"),
]

# Generate 3D assets for all inputs in parallel

meshes_batch = pipeline.run(
    image=images,               # List of PIL Images

    num_samples=1,              # One 3D output per input image

    batch_size=len(images),     # Pass to samplers via sampler_params

    preprocess_image=True,
)

# Save results

for i, mesh in enumerate(meshes_batch):
    mesh.save(f"output_{i}.glb")

The batch_size parameter propagates through to the internal samplers (sample_sparse_structure, sample_shape_slat, sample_tex_slat), which allocate noise tensors of shape (batch, ...) accordingly.

Method 2: Sequential Processing (Fallback)

If memory constraints prevent full batching, process images sequentially while reusing the pipeline instance. This approach uses the standard single-image API but iterates externally.

meshes = []
for img in images:
    result = pipeline.run(
        image=img,
        num_samples=1,
        preprocess_image=True,
    )
    meshes.append(result[0])  # Each call returns a list of MeshWithVoxel objects

How Batch Processing Works Under the Hood

The TRELLIS.2 pipeline supports batch generation through three architectural components:

1. Batched Conditioning in image_cond_model

The feature extractor in trellis2/modules/image_feature_extractor.py accepts either a batched tensor or a list of PIL images. It stacks inputs into a tensor of shape (B, C, H, W) and returns global features of shape (B, N, D).

# From trellis2/pipelines/trellis2_image_to_3d.py

cond = self.image_cond_model(image)  # Handles batch dimension automatically

2. Sampler Batch Allocation

All flow-matching samplers respect the batch_size argument passed via sampler_params. For example, sample_sparse_structure allocates noise using the specified batch size:

# Sparse structure sampling allocates batch-sized noise

noise = torch.randn(num_samples, in_channels, reso, reso, reso).to(self.device)

3. SparseTensor Batch Dimension

The SparseTensor class defined in trellis2/modules/sparse.py carries the batch dimension through the shape-SLat and tex-SLat stages. The decoders (decode_shape_slat, decode_tex_slat) iterate over this dimension automatically when processing the final latent representation.

Pipeline Step	Function	Batch Handling	Source File
Pre-processing	preprocess_image	Applied per-image or batched	trellis2/pipelines/trellis2_image_to_3d.py (lines 27-33)
Feature Extraction	get_cond	Returns (B, N, D) tensor	trellis2/pipelines/trellis2_image_to_3d.py (lines 74-78)
Structure Sampling	sample_sparse_structure	Noise shape (batch, C, R, R, R)	trellis2/pipelines/trellis2_image_to_3d.py
Shape Latent Sampling	sample_shape_slat	SparseTensor with batched feats	trellis2/pipelines/trellis2_image_to_3d.py (lines 53-56)
Texture Latent Sampling	sample_tex_slat	Concatenates batched conditioning	trellis2/pipelines/trellis2_image_to_3d.py (lines 11-13)
Mesh Decoding	decode_latent	Loops over batch dimension	trellis2/pipelines/trellis2_image_to_3d.py (lines 70-72)

Key Source Files for Batch Implementation

Reference these files when customizing batch behavior or debugging memory issues:

• trellis2/pipelines/trellis2_image_to_3d.py – Main inference pipeline containing the run method and sampler orchestration.
• trellis2/modules/image_feature_extractor.py – Implements DinoV2FeatureExtractor and DinoV3FeatureExtractor with native batch support.
• trellis2/modules/sparse.py – Defines SparseTensor, the data structure that maintains batch dimensions through sparse convolution operations.
• trellis2/trainers/flow_matching/flow_matching.py – Generic flow-matching sampler implementation used by structure and latent sampling stages.
• trellis2/trainers/utils.py – Helper utilities including recursive_to_device for moving batched tensors between CPU and GPU.

Summary

• Pass a list of PIL.Image objects to pipeline.run() instead of a single image to trigger batch processing.
• Set batch_size=len(images) in the sampler parameters to ensure noise tensors allocate correctly for your input count.
• The image_cond_model automatically handles batching, returning features of shape (B, N, D).
• SparseTensor objects maintain the batch dimension throughout shape-SLat and tex-SLat generation.
• For memory-constrained environments, fall back to external iteration while keeping the pipeline instance loaded.

Frequently Asked Questions

Does TRELLIS.2 natively support multiple input images per call?

Yes, the architecture supports batch generation natively. The image_cond_model in trellis2/modules/image_feature_extractor.py accepts lists of PIL images or batched tensors, and all samplers accept a batch_size parameter. However, you must explicitly pass the batch_size or a list of images to activate this behavior, as the default run configuration targets single-image generation.

How does memory usage scale with batch size?

Memory consumption grows linearly with batch size during the sparse structure sampling and latent decoding phases, as these stages allocate dense tensors of shape (batch, channels, resolution, resolution, resolution). Shape-SLat and tex-SLat sampling use sparse representations (SparseTensors), which partially mitigate memory growth for sparse geometries, but the decoders process each batch element to produce final meshes.

Can I generate multiple variants per image in a batch?

Yes. Set num_samples to the desired number of 3D variants per input image. The total number of outputs will be batch_size × num_samples. Ensure your GPU memory can accommodate the combined noise tensors, as each sample requires independent random initialization in the samplers.

What input formats are supported for batch processing?

The pipeline accepts a list of PIL Image objects (RGB or RGBA) or a pre-batched torch.Tensor of shape (B, C, H, W). If using tensors, disable preprocess_image (set to False) and ensure images are resized to 512×512 or 1024×1024 to match the conditioning model's expected input resolutions. Mixed sizes in a single batch are not supported without individual pre-processing.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}