Research

AI-generated text

Fine-tuning TRELLIS.2 with LoRA adapters on footwear scans

Researchers provide scripts and container images to fine-tune TRELLIS.2 for a specific object category (footwear) using LoRA adapters and the Google Scanned Objects (GSO) dataset.

Fine-tuning TRELLIS.2 with LoRA adapters on footwear scans

This article describes how TRELLIS.2 was fine-tuned for footwear generation using LoRA (low-rank adaptation) adapters. The authors supply Docker container images and scripts that run locally or on Hugging Face Jobs. The experiment used 188 footwear objects from the Google Scanned Objects (GSO) dataset: 179 objects for training and 9 held out for evaluation.

Why this matters

General-purpose 3D models such as TRELLIS.2 perform well across many categories but can make errors in details that matter for specific applications (for example, inaccurate soles on shoes). The objective here is to improve geometry for a target category without overwriting the model’s existing useful priors.

What they fine-tuned and why

TRELLIS.2’s pipeline has five stages: background removal (BiRefNet), image encoding (DINOv3-L), coarse structure prediction (sparse-structure flow DiT), detailed geometry generation (shape flow DiT), and material/texture generation (texture flow DiT). The team inserted LoRA adapters into the two geometry-controlling DiTs (sparse-structure and shape) and froze the texture flow DiT.

Rationale:

  • The structure and shape flows directly control geometry, so adapting them targets the problem (shoe soles).
  • The available dataset contains only diffuse color maps with baked lighting, not separate material maps. Training the texture flow DiT on these would risk learning lighting as surface properties and cause catastrophic forgetting.

LoRA configuration and resources

  • LoRA adapters are added via Hugging Face’s PEFT tools, with custom injection due to TRELLIS.2’s nonstandard forward pass.
  • The configuration trains roughly 9.3 million adapter parameters in total.
  • Adapter ranks: r=8 for attention and MLP layers in the sparse-structure DiT; r=16 for attention layers in the shape DiT.
  • Hardware: runs used A100 80 GB GPUs on Hugging Face Jobs.

Data preparation and execution options

  • Data: a 188-object footwear subset from GSO (a subset of the 1,030-object GSO collection). For each object, Blender renders 16 views; preprocessing for the 188 objects took about 1 hour on an a100-large (80 GB) job with the provided pipeline.
  • Two container images are provided: one for preprocessing (renders + conversion to TRELLIS.2 latents) and one for LoRA fine-tuning.
  • The workflow can be run locally via Docker or remotely via Hugging Face Jobs. Example commands show passing HF_TOKEN as a secret, selecting --flavor a100-large, timeouts, and other parameters.

Internal representations

  • TRELLIS.2 uses a unified latent representation; its decoders can emit various 3D output formats. O-Voxel is TRELLIS.2’s sparse voxel format storing only surface voxels, which makes CPU/GPU conversions fast.
  • The sparse flows operate on roughly 10% of voxels, reducing wasted compute on empty space.

Fine-tuning details

  • Preprocessing: frozen TRELLIS.2 encoders convert meshes to latents which are saved for reuse. For training, one random view out of 16 is used as conditioning per sample.
  • Dataset split: 179 training objects, 9 held-out for evaluation.
  • LoRA injection: because TRELLIS.2 uses a custom forward, inject_adapter_in_model() is used instead of the standard PeftModel wrapper. Trainable adapter weights are kept in fp32 while bf16 autocasting is used for forward passes.
  • Optimization: AdamW with gradient accumulation (4), gradient checkpointing, gradient clipping, and adapter checkpoints saved every 500 steps.
  • Loss: rectified flow-matching, matching the pretraining objective of TRELLIS.2.

Training loop summary:

  • For each batch: sample timestep t (logitNormal for sparse-structure DiT, uniform for shape DiT), mix clean latents with noise according to t, compute model prediction of velocity, and measure MSE against the target velocity.
  • DINOv3 encodes conditioning images for classifier-free guidance; a random subset of samples are unconditioned (p_uncond).

Run details and results

  • Each LoRA run used one A100 80 GB GPU on Hugging Face Jobs.
  • Run times: sparse-structure DiT run took 7.2 hours; shape DiT run took 3.6 hours.
  • Quantitative results on the nine held-out objects (means):
    • F-score @ 0.02: base model 0.276 → with LoRAs 0.803
    • Chamfer distance: 0.174 → 0.033 (lower is better)
    • Normal consistency: 0.530 → 0.804 (higher is better)
  • These metrics show substantial improvement in geometry quality from the LoRA adapters. Visual comparisons on selected held-out examples indicate that soles generated by the fine-tuned model are closer to the scans.

Code, artifacts and reuse

  • The preprocessed latents are saved so multiple training runs can reuse them.
  • The repository suvadityamuk/trellis-2-finetune-gso contains training code, LoRA configs and the trained adapters. The authors also provide commands and container references for reproducing the preprocessing and training steps.

Conclusions and next steps

  • Adding LoRA adapters to the geometry-related DiTs of TRELLIS.2 notably improved footwear geometry while keeping the rest of the model frozen.
  • Key lesson: choose the training targets that align with both the problem and the available data; freezing the texture stage here was important because of the dataset’s limitations.
  • Future work could test the same approach on other 3D generation models (for example, Pixal3D, UniLat3D) or on models using different representations (Gaussian splatting, world-generation models), which may require different pipeline choices.

The provided containers and scripts make the pipeline reproducible locally or on Hugging Face Jobs for teams that want to adapt TRELLIS.2 to other object categories.