How pbr_attr_layout Maps Base Color, Metallic, Roughness, and Alpha to Mesh Attributes in TRELLIS.2

The pbr_attr_layout dictionary in TRELLIS.2 maps a dense six-channel texture tensor to material properties using Python slice objects, allocating channels 0–2 to base color RGB, channel 3 to metallic, channel 4 to roughness, and channel 5 to alpha.

In the TRELLIS.2 framework for texturing and image-to-3D generation, material attributes are packed into a single six-channel tensor to optimize memory and computation. The pbr_attr_layout configuration defines how these channels unpack into standard PBR (Physically Based Rendering) mesh attributes during post-processing.

Understanding the Six-Channel PBR Tensor

TRELLIS.2 represents PBR material data as a dense 3-D tensor with shape (H, W, 6), where the final dimension contains all material properties concatenated together. This design allows the model to predict complete material information in a single forward pass without separate texture heads for each attribute.

The channel dimension follows a strict ordering: RGB base color occupies the first three indices, while single-channel properties (metallic, roughness, and alpha) occupy the remaining three indices sequentially.

How pbr_attr_layout Defines Channel Mappings

The pbr_attr_layout dictionary is instantiated in the pipeline constructor using Python slice objects to index specific channel ranges. As implemented in trellis2/pipelines/trellis2_texturing.py at lines 60–65, the layout maps each material property to its corresponding tensor indices:

self.pbr_attr_layout = {
    'base_color': slice(0, 3),  # Channels 0, 1, 2

    'metallic':   slice(3, 4),  # Channel 3

    'roughness':  slice(4, 5),  # Channel 4

    'alpha':      slice(5, 6),  # Channel 5

}

Base Color (RGB Channels)

The base_color property uses slice(0, 3) to extract channels 0 through 2, representing the red, green, and blue color values of the material surface. Because slice notation is half-open, this captures exactly three channels suitable for RGB texture data.

Metallic, Roughness, and Alpha Channels

The remaining properties use single-channel slices:

• metallic: slice(3, 4) isolates channel 3 for the metallic coefficient (0 = dielectric, 1 = metallic)
• roughness: slice(4, 5) isolates channel 4 for surface roughness values
• alpha: slice(5, 6) isolates channel 5 for transparency/opacity masks

Attribute Extraction in postprocess_mesh

During mesh generation, the postprocess_mesh method utilizes pbr_attr_layout to unpack the tensor into separate numpy arrays for further processing. At lines 336–339 of trellis2/pipelines/trellis2_texturing.py, the extraction follows this pattern:

# attrs shape: (H, W, 6)

base_color = attrs[..., self.pbr_attr_layout['base_color']].cpu().numpy()
metallic   = attrs[..., self.pbr_attr_layout['metallic']].cpu().numpy()
roughness  = attrs[..., self.pbr_attr_layout['roughness']].cpu().numpy()
alpha      = attrs[..., self.pbr_attr_layout['alpha']].cpu().numpy()

These extracted arrays undergo post-processing operations including scaling to 0–255 ranges, inpainting missing pixels, and conversion to trimesh.visual.material.PBRMaterial objects that attach directly to the output mesh.

Cross-Pipeline Consistency

To ensure identical behavior across different entry points, the image-to-3D pipeline replicates this exact layout definition. In trellis2/pipelines/trellis2_image_to_3d.py at lines 73–78, the same dictionary structure initializes the pbr_attr_layout attribute, guaranteeing that models loaded through either the texturing or image-to-3D pipelines interpret channel data identically.

The renderer components in trellis2/renderers/pbr_mesh_renderer.py expect this specific channel ordering when processing meshes, creating a consistent contract between tensor generation and visualization throughout the codebase.

Practical Code Example: Splitting PBR Tensors

When working with raw model outputs, you can use the layout to decompose the six-channel tensor into individual material textures:

import torch
import numpy as np
from PIL import Image

def extract_pbr_textures(attrs_tensor, layout):
    """
    Args:
        attrs_tensor: torch.Tensor of shape (H, W, 6)
        layout: dict with slice objects from pbr_attr_layout
    Returns:
        Tuple of (base_color, metallic, roughness, alpha) as numpy arrays
    """
    base = attrs_tensor[..., layout['base_color']].cpu().numpy()
    metallic = attrs_tensor[..., layout['metallic']].cpu().numpy()
    roughness = attrs_tensor[..., layout['roughness']].cpu().numpy()
    alpha = attrs_tensor[..., layout['alpha']].cpu().numpy()
    
    return base, metallic, roughness, alpha

# Usage with TRELLIS.2 pipeline

layout = pipeline.pbr_attr_layout
base, met, rough, alpha = extract_pbr_textures(model_output, layout)

# Convert to 8-bit textures for export

base_img = Image.fromarray((base * 255).astype(np.uint8))

Summary

• pbr_attr_layout uses slice objects to map channels 0–2 to base_color RGB, channel 3 to metallic, channel 4 to roughness, and channel 5 to alpha
• The layout is defined in trellis2_texturing.py (lines 60–65) and trellis2_image_to_3d.py (lines 73–78)
• Extraction occurs in postprocess_mesh at lines 336–339 via tensor slicing with the layout dictionary
• This six-channel packing enables efficient single-pass prediction while maintaining compatibility with standard glTF PBR material workflows

Frequently Asked Questions

What is the exact channel order in the TRELLIS.2 PBR tensor?

The channel order is strictly RGB-base-color, metallic, roughness, alpha, occupying indices 0 through 5 respectively. Channels 0–3 hold the base color, while channels 3, 4, and 5 contain scalar values for metallic, roughness, and alpha transparency.

Where is pbr_attr_layout initialized in the codebase?

The layout is initialized in the constructor of Trellis2TexturingPipeline located at trellis2/pipelines/trellis2_texturing.py lines 60–65, and identically in Trellis2ImageTo3DPipeline at trellis2/pipelines/trellis2_image_to_3d.py lines 73–78.

How does TRELLIS.2 handle the alpha channel during mesh export?

The alpha channel extracted via pbr_attr_layout['alpha'] (slice 5–6) is combined with the base color RGB to create an RGBA texture, then assigned to the baseColorTexture parameter of the PBRMaterial. The alphaMode is typically set to 'OPAQUE' or 'BLEND' depending on the presence of transparent pixels.

Can I modify pbr_attr_layout to support additional material channels?

While technically possible by adjusting the slice definitions, modifying pbr_attr_layout requires corresponding changes to the model architecture (output channel dimension), the renderer expectations in pbr_mesh_renderer.py, and the voxel post-processing utilities. The current six-channel layout is hardcoded across multiple components to match the trained model weights.

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