DOCSEARCHbyAlgolia
Submit Docs

UniRig: Unified Framework for Automatic 3D Model Rigging

repository·main·Indexed 23 days ago

https://github.com/vast-ai-research/unirig

UniRig is a unified framework for automatic 3D model rigging that predicts skeletal structures and skinning weights for diverse 3D assets. It utilizes a GPT-like transformer for autoregressive skeleton prediction via Skeleton Tree Tokenization and a Bone-Point Cross Attention mechanism for per-vertex skinning weight and attribute prediction. The framework supports .obj, .fbx, .glb, and .vrm input formats and includes tools for training, inference, and merging predicted results with original meshes.

Tokens
7.2K
Snippets
11
Records
36
Agent score
82%

What's inside UniRig

  1. How UniRig works: Skeleton and Skinning stages

    main

    UniRig is a unified framework for automatic 3D model rigging consisting of two main stages:

    1. Skeleton Prediction: Uses a GPT-like transformer to autoregressively predict a topologically valid skeleton hierarchy via Skeleton Tree Tokenization.
    2. Skinning Weight & Attribute Prediction: Uses a Bone-Point Cross Attention mechanism to predict per-vertex skinning weights and bone attributes (like stiffness) based on the predicted skeleton and the input mesh geometry.

    Current Status: Skeleton prediction and skinning weight prediction are currently available. Bone attribute prediction is planned for a future release.

  2. Prepare custom data for training

    main

    To use a custom dataset, you must first configure the data paths in configs/data/rignet.yaml. Set input_dataset_dir to your original model folder and output_dataset_dir to the destination for processed .npz files. Once configured, run the preprocessing script to generate the required dataset format.

    Note: The dataloader expects processed data to be located under <output_dataset_dir>/dataset_clean/<relative path in datalist>/raw_data.npz.

    bash launch/inference/preprocess.sh --config configs/data/<yourdata> --num_runs <number of threads to run>

  3. Perform Skeleton Inference (Prediction)

    main

    To run inference on a trained skeleton model, create a new task configuration file (e.g., configs/task/rignet_ar_inference_scratch.yaml) with mode: predict and point resume_from_checkpoint to your trained .ckpt file.

    Required Configuration for Skeleton Inference:

    • mode: predict
    • components.tokenizer: Must match the tokenizer used during training.
    • components.model: Must match the model used during training.
    • writer.export_npz: predict_skeleton
    • writer.export_obj: skeleton
    • writer.export_fbx: skeleton
    • trainer.inference_mode: (Implicitly handled by mode: predict)

    Run the inference using the generate_skeleton.sh script.

    bash launch/inference/generate_skeleton.sh --input examples/giraffe.glb --output examples/giraffe_skeleton.fbx --skeleton_task configs/task/rignet_ar_inference_scratch.yaml

  4. Train the Skeleton Model

    main

    Training the skeleton model requires a complete configuration set including data, transform, tokenizer, system, model, and task.

    Key Configuration Notes:

    • Model: In configs/model/unirig_rignet.yaml, ensure n_positions is greater than the sum of the conditional embedding length and the maximum number of skeleton tokens.
    • System: configs/system/ar_train_rignet.yaml controls the training process, including generation export intervals and sampling methods.
    • Task: configs/task/train_rignet_ar.yaml integrates all components and configures loss, optimizer, and scheduler. It also contains the trainer section for GPU/node usage.
    • Logging: You can comment out the wandb and checkpoint sections in the task config if you do not require Weights & Biases logging or model checkpoints.

    Checkpoints are saved to experiments/<experimentname>.

    python run.py --task=configs/task/train_rignet_ar.yaml

  5. Perform Skin Inference (Prediction)

    main

    To run inference on a trained skin model, create a task configuration (e.g., configs/task/rignet_skin_inference_scratch.yaml) with mode: predict and the path to your checkpoint in resume_from_checkpoint.

    Required Configuration for Skin Inference:

    • mode: predict
    • components.system: skin
    • components.transform: inference_skin_transform (Note: does not need skin vertex groups)
    • components.model: Must match the model used during training.
    • writer.export_fbx: result_fbx
    • trainer.inference_mode: True

    Run the inference using the generate_skin.sh script.

    bash launch/inference/generate_skin.sh --input examples/skeleton/giraffe.fbx --output results/giraffe_skin.fbx --skin_task configs/task/rignet_skin_inference_scratch.yaml

  6. Install UniRig

    main

    Follow these steps to set up the UniRig environment.

    Prerequisites:

    • Python 3.11
    • PyTorch (version >= 2.3.1)

    Setup Steps:

    1. Clone the repository:
      git clone https://github.com/VAST-AI-Research/UniRig
cd UniRig
    2. Create and activate a virtual environment:
      conda create -n UniRig python=3.11
conda activate UniRig
    3. Install core dependencies:
      python -m pip install torch torchvision
python -m pip install -r requirements.txt
python -m pip install numpy==1.26.4
    4. Install specialized dependencies (replace placeholders with your specific versions):
      • spconv: python -m pip install spconv-{you-cuda-version}
      • torch_scatter and torch_cluster: Install from PyG wheels using python -m pip install torch_scatter torch_cluster -f <URL> --no-cache-dir.

    Note on flash_attn: If you encounter installation errors for flash_attn, follow the guide in its original repository.

    git clone https://github.com/VAST-AI-Research/UniRig
cd UniRig
conda create -n UniRig python=3.11
conda activate UniRig
python -m pip install torch torchvision
python -m pip install -r requirements.txt
python -m pip install spconv-{you-cuda-version}
python -m pip install torch_scatter torch_cluster -f https://data.pyg.org/whl/torch-{your-torch-version}+{your-cuda-version}.html --no-cache-dir
python -m pip install numpy==1.26.4

  7. Train the Skin Model

    main

    Skinning training follows a similar process to skeleton training. Use configs/task/train_rignet_skin.yaml to initiate training.

    Memory Management: This task is memory-intensive (requires ~60GB VRAM for default settings). If you encounter Out-of-Memory (OOM) errors, you can:

    1. Set batch_size: 1 in the data config.
    2. Increase accumulate_grad_batches in the task config.
    3. Decrease num_train_vertex in configs/model/unirig_skin.yaml.

    Troubleshooting: If you encounter pyrender issues, change the backend in configs/transform/train_rignet_skin_transform by setting vertex_group_confis/kwargs/voxel_skin/backend to open3d. Note that this change must also be applied when running in prediction mode.

    python run.py --task=configs/task/train_rignet_skin.yaml

  8. How FrequencyPositionalEmbedding works

    main

    The FrequencyPositionalEmbedding module implements a sin/cosine positional embedding to enrich input features with multi-frequency information. This is commonly used to help neural networks capture fine-grained spatial details.

    Given an input x of shape [n_batch, ..., c_dim], it transforms each dimension into a vector containing sine and cosine components at multiple frequencies.

    Parameters:

    • num_freqs (int): Number of frequencies to use (default: 6).
    • logspace (bool): If True, frequencies are spaced exponentially: [2^(0/num_freqs), ..., 2^(i/num_freqs), ...]. If False, they are linearly spaced between [1.0, 2^(num_freqs - 1)].
    • input_dim (int): The dimensionality of the input (default: 3).
    • include_input (bool): If True, the original input x is concatenated to the output (default: True).
    • include_pi (bool): If True, frequencies are multiplied by $\pi$ (default: True).

    Output Dimension:

    • If include_input is True: input_dim * (num_freqs * 2 + 1)
    • If include_input is False: input_dim * num_freqs * 2

  9. Use EventStorage for metric logging

    main

    The EventStorage class is the primary interface for storing and managing training metrics (scalars, images, histograms). It is designed to be used as a context manager to enable global access via get_event_storage().

    To use it, wrap your training loop or relevant code block in a with EventStorage(...) statement. This allows any component in your code to retrieve the active storage instance and log data without passing the object explicitly through every function call.

  10. Visualize RigXL data as FBX

    main

    You can export data from a RigXL .npz file to an FBX model using the RawData class.

    from src.data.raw_data import RawData
raw_data = RawData.load("dataset_clean/rigxl/12345/raw_data.npz")
raw_data.export_fbx("res.fbx")

  11. Configure UniRigSkin initialization parameters

    main

    When initializing UniRigSkin, you must provide a configuration dictionary via **kwargs. The following keys are used to define the model architecture:

    KeyDescription
    num_train_vertexNumber of vertices to sample during training
    feat_dimFeature dimension for the model
    num_headsNumber of attention heads
    grid_sizeGrid size for mesh encoding
    mlp_dimDimension of the MLP layers
    num_bone_attnNumber of attention blocks in the BoneEncoder
    num_mesh_bone_attnNumber of cross-attention blocks between mesh and bones
    bone_embed_dimEmbedding dimension for bones
    voxel_maskPower used for skin mask normalization (default: 2)

    Additionally, you must provide two configuration dictionaries:

    • mesh_encoder: Configuration for the local mesh encoder (e.g., ptv3obj).
    • global_encoder: Configuration for the global mesh encoder (e.g., michelangelo_encoder).