Motion VQ-VAE model update

.gitattributes

@@ -1,35 +1,5 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
 *.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,224 @@
+---
+license: cc-by-nc-4.0
+tags:
+  - motion
+  - vqvae
+  - motion-tokenization
+  - motion-generation
+  - human-motion
+  - vector-quantization
+language:
+  - en
+library_name: transformers
+pipeline_tag: feature-extraction
+datasets:
+  - MotionMillion
+---
+
+# MotionVQVAE
+
+A Multi-Group Vector Quantized VAE (MG-VQVAE) trained on the MotionMillion dataset for motion tokenization and reconstruction.
+
+> ⚠️ **License Notice**: This model is released under **CC BY-NC 4.0** (Creative Commons Attribution-NonCommercial 4.0). The training data includes datasets with mixed licensing terms, some of which restrict commercial use. **This model is for research and non-commercial use only.**
+
+> 📋 **Body Model**: This model was trained on motion data using the **SMPL body model** (22 joints). Input motions must be in SMPL skeleton format.
+
+## Model Description
+
+MotionVQVAE learns to compress human motion sequences into discrete tokens using a Multi-Group Vector Quantization approach. The model can:
+
+- **Tokenize** motion sequences into discrete tokens for downstream generation tasks
+- **Reconstruct** motions from tokens with high fidelity
+- **Compress** variable-length motions with 4× temporal downsampling
+
+### Multi-Group VQ Architecture
+
+Instead of a single codebook, MG-VQVAE uses **64 parallel groups**, each with its own 512-code codebook. This provides:
+- Effective codebook size: $512^{64} \approx 2.47 \times 10^{173}$ combinations
+- Fine-grained control over different motion aspects
+- Better reconstruction quality through distributed quantization
+
+## Usage
+
+### Installation
+
+```bash
+pip install torch huggingface_hub numpy
+```
+
+### Download the Model Code
+
+Download `motion_vqvae_hf.py` from this repository or copy it to your project.
+
+### Quick Start
+
+```python
+from motion_vqvae_hf import MotionVQVAE
+import numpy as np
+
+# Load model (auto-downloads from HuggingFace)
+model = MotionVQVAE.from_pretrained("khania/motion-vqvae")
+
+# Prepare motion data (272-dim absolute root format)
+motion = np.random.randn(120, 272).astype(np.float32)  # Replace with real motion
+
+# Encode motion to tokens
+tokens = model.encode(motion)  # Returns token indices for each group
+print(f"Tokens shape: {tokens.shape}")  # (64, 1, 30) - 64 groups, batch=1, T/4 timesteps
+
+# Decode tokens back to motion
+motion_recon = model.decode(tokens)
+print(f"Reconstructed motion shape: {motion_recon.shape}")  # (1, 120, 272)
+
+# Full forward pass (encode + decode)
+motion_recon, tokens = model(motion)
+```
+
+### Batch Processing
+
+```python
+# Process multiple motions
+motions = [
+    np.random.randn(100, 272).astype(np.float32),
+    np.random.randn(150, 272).astype(np.float32),
+    np.random.randn(80, 272).astype(np.float32),
+]
+
+# Encode batch (will pad to max length)
+tokens = model.encode_batch(motions)
+
+# Decode batch
+motions_recon = model.decode_batch(tokens)
+```
+
+### Access Codebook
+
+```python
+# Get quantized embeddings for analysis
+embeddings = model.get_codebook_embeddings()
+print(f"Codebook shape: {embeddings.shape}")  # (64, 512, 8) - 64 groups, 512 codes, 8-dim each
+```
+
+## Model Architecture
+
+| Component | Details |
+|-----------|---------|
+| **Encoder** | 1D CNN with residual blocks |
+| **Decoder** | 1D CNN with residual blocks |
+| **Width** | 1024 |
+| **Depth** | 3 residual blocks per stage |
+| **Downsampling** | 4× (stride 2, 2 stages) |
+| **Quantizer** | Multi-Group VQ with EMA updates |
+| **Groups** | 64 |
+| **Codebook Size** | 512 codes per group |
+| **Code Dimension** | 8 per group (512 total) |
+| **Total Parameters** | ~73M |
+
+## Motion Format
+
+The model expects **272-dimensional motion features in absolute root format** based on the **SMPL body model** (22 joints).
+
+### SMPL Body Model Requirement
+
+This model was trained exclusively on motion data represented using the [SMPL body model](https://smpl.is.tue.mpg.de/). Your input motions must:
+
+- Use the **SMPL skeleton** with 22 joints
+- Follow the SMPL joint ordering
+- Be converted to the 272-dimensional HumanML3D-style representation
+
+If your motion data uses a different skeleton (e.g., CMU, Mixamo, custom rigs), you must first retarget it to SMPL before using this model.
+
+### Feature Dimensions
+
+| Dimensions | Description |
+|------------|-------------|
+| `[0:2]` | Root XZ velocities |
+| `[2:8]` | Absolute heading rotation (6D representation) |
+| `[8:74]` | Local joint positions (22 joints × 3) |
+| `[74:140]` | Local joint velocities (22 joints × 3) |
+| `[140:272]` | Joint rotations in 6D (22 joints × 6) |
+
+The model automatically normalizes input motions using the bundled mean/std statistics.
+
+## Training Details
+
+| Parameter | Value |
+|-----------|-------|
+| **Dataset** | MotionMillion |
+| **Batch Size** | 128 |
+| **Training Iterations** | 300,000 |
+| **Learning Rate** | 2e-4 |
+| **LR Schedule** | Step decay at 50K, 400K |
+| **Loss Function** | L1 Smooth + Commitment |
+| **Commitment Weight** | 0.02 |
+| **Window Size** | 64 frames |
+
+### Loss Weights
+
+| Component | Weight |
+|-----------|--------|
+| Root XZ Velocity | 3.0 |
+| Root Rotation | 1.5 |
+| Joint Position | 0.1 |
+| Joint Velocity | 0.5 |
+| Joint Rotation | 5.0 |
+| Velocity Temporal | 0.5 |
+
+## Performance
+
+Final evaluation metrics at 300K iterations:
+
+| Metric | Value |
+|--------|-------|
+| **Reconstruction Loss** | 0.0095 |
+| **Commitment Loss** | 0.0255 |
+| **Perplexity** | 508.57 |
+| **Codebook Utilization** | 100% |
+
+### Per-Component Reconstruction Loss (Eval)
+
+| Component | Loss |
+|-----------|------|
+| Root XZ Velocity | 0.00107 |
+| Root Rotation | 0.00029 |
+| Joint Position | 0.00851 |
+| Joint Velocity | 0.02301 |
+| Joint Rotation | 0.00383 |
+
+## Files in This Repository
+
+| File | Size | Description |
+|------|------|-------------|
+| `config.json` | ~300 B | Model configuration |
+| `pytorch_model.bin` | ~280 MB | Model weights (~73M parameters) |
+| `mean.npy` | 1.2 KB | Motion normalization mean (272,) |
+| `std.npy` | 1.2 KB | Motion normalization std (272,) |
+| `motion_vqvae_hf.py` | ~20 KB | Model implementation |
+
+## Use Cases
+
+- **Motion Generation**: Tokenize motions for autoregressive or diffusion-based generation
+- **Motion Compression**: Efficiently store motion data as discrete tokens
+- **Motion Editing**: Manipulate tokens for motion modification
+- **Downstream Tasks**: Use tokens as input for text-to-motion models
+
+## Limitations
+
+- Trained on English text descriptions only (for associated metadata)
+- Motion format is specific to HumanML3D-style 272-dim representation
+- 4× temporal downsampling may lose very fine-grained details
+- Best performance on motions similar to training distribution (daily activities, sports, etc.)
+
+## Citation
+
+```bibtex
+@article{motionmillion2026,
+  title={MotionMillion: A Large-Scale Motion-Language Dataset},
+  author={...},
+  year={2026}
+}
+```
+
+## License
+
+**CC BY-NC 4.0** (Creative Commons Attribution-NonCommercial 4.0 International)

config.json

@@ -0,0 +1,15 @@
+{
+  "input_dim": 272,
+  "code_dim": 512,
+  "nb_code": 512,
+  "num_groups": 64,
+  "down_t": 2,
+  "stride_t": 2,
+  "width": 1024,
+  "depth": 3,
+  "dilation_growth_rate": 3,
+  "kernel_size": 3,
+  "activation": "relu",
+  "use_mgvq": true,
+  "absolute_root": true
+}

mean.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3e3ce8012ec7085209c805c3d9f8deb56bc447e8901b8f30fea8da6a841f302
+size 1216

motion_vqvae_hf.py

@@ -0,0 +1,760 @@
+"""
+MotionVQVAE - Motion Vector Quantized VAE for HuggingFace
+
+Load and use the MotionVQVAE model for motion tokenization and reconstruction.
+
+Usage:
+    from motion_vqvae_hf import MotionVQVAE
+    
+    # Load from HuggingFace Hub
+    model = MotionVQVAE.from_pretrained("khania/motion-vqvae")
+    
+    # Encode motion to tokens
+    tokens = model.encode(motion_array)  # (B, T, 272) -> (num_groups, B, T')
+    
+    # Decode tokens back to motion
+    motion_recon = model.decode(tokens)  # (num_groups, B, T') -> (B, T, 272)
+    
+    # Full forward pass
+    motion_recon, tokens = model(motion_array)
+"""
+
+import os
+import json
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from typing import List, Union, Optional, Dict, Any, Tuple
+from pathlib import Path
+
+try:
+    from huggingface_hub import snapshot_download
+    HF_HUB_AVAILABLE = True
+except ImportError:
+    HF_HUB_AVAILABLE = False
+
+
+# =============================================================================
+# Encoder / Decoder Components (matching original architecture exactly)
+# =============================================================================
+
+class ResConv1DBlock(nn.Module):
+    """Residual 1D Convolution Block - matches original models/resnet.py exactly."""
+    
+    def __init__(self, n_in: int, n_state: int, dilation: int = 1, 
+                 activation: str = 'relu', norm: str = None, kernel_size: int = 3):
+        super().__init__()
+        padding = dilation * (kernel_size - 1) // 2
+        self.norm = norm
+        
+        # Norm layers
+        if norm == "LN":
+            self.norm1 = nn.LayerNorm(n_in)
+            self.norm2 = nn.LayerNorm(n_in)
+        elif norm == "GN":
+            self.norm1 = nn.GroupNorm(num_groups=32, num_channels=n_in, eps=1e-6, affine=True)
+            self.norm2 = nn.GroupNorm(num_groups=32, num_channels=n_in, eps=1e-6, affine=True)
+        elif norm == "BN":
+            self.norm1 = nn.BatchNorm1d(num_features=n_in, eps=1e-6, affine=True)
+            self.norm2 = nn.BatchNorm1d(num_features=n_in, eps=1e-6, affine=True)
+        else:
+            self.norm1 = nn.Identity()
+            self.norm2 = nn.Identity()
+        
+        # Activation layers
+        if activation == "relu":
+            self.activation1 = nn.ReLU()
+            self.activation2 = nn.ReLU()
+        elif activation == "silu":
+            self.activation1 = nn.SiLU()
+            self.activation2 = nn.SiLU()
+        elif activation == "gelu":
+            self.activation1 = nn.GELU()
+            self.activation2 = nn.GELU()
+        else:
+            self.activation1 = nn.ReLU()
+            self.activation2 = nn.ReLU()
+        
+        # Convolution layers - MUST be named conv1 and conv2 to match checkpoint
+        self.conv1 = nn.Conv1d(n_in, n_state, kernel_size, 1, padding, dilation)
+        self.conv2 = nn.Conv1d(n_state, n_in, 1, 1, 0)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_orig = x
+        if self.norm == "LN":
+            x = self.norm1(x.transpose(-2, -1))
+            x = self.activation1(x.transpose(-2, -1))
+        else:
+            x = self.norm1(x)
+            x = self.activation1(x)
+        
+        x = self.conv1(x)
+        
+        if self.norm == "LN":
+            x = self.norm2(x.transpose(-2, -1))
+            x = self.activation2(x.transpose(-2, -1))
+        else:
+            x = self.norm2(x)
+            x = self.activation2(x)
+        
+        x = self.conv2(x)
+        x = x + x_orig
+        return x
+
+
+class Resnet1D(nn.Module):
+    """1D Residual Network - matches original models/resnet.py exactly.
+    
+    Uses self.model = nn.Sequential(*blocks) to match checkpoint key structure.
+    """
+    
+    def __init__(self, n_in: int, n_depth: int, dilation_growth_rate: int = 1,
+                 reverse_dilation: bool = False, activation: str = 'relu', 
+                 norm: str = None, kernel_size: int = 3):
+        super().__init__()
+        
+        blocks = [
+            ResConv1DBlock(n_in, n_in, dilation=dilation_growth_rate ** depth, 
+                          activation=activation, norm=norm, kernel_size=kernel_size)
+            for depth in range(n_depth)
+        ]
+        if reverse_dilation:
+            blocks = blocks[::-1]
+        
+        # MUST be named 'model' to match checkpoint keys like 'model.0.conv1.weight'
+        self.model = nn.Sequential(*blocks)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+
+class Encoder(nn.Module):
+    """1D CNN Encoder - matches original models/encdec.py exactly.
+    
+    Uses self.model = nn.Sequential(*blocks) to match checkpoint key structure:
+    - model.0: Conv1d (input projection)
+    - model.1: ReLU
+    - model.2: Sequential(Conv1d, Resnet1D) for first downsample
+    - model.3: Sequential(Conv1d, Resnet1D) for second downsample
+    - model.4: Conv1d (output projection)
+    """
+    
+    def __init__(
+        self,
+        input_emb_width: int = 272,
+        output_emb_width: int = 512,
+        down_t: int = 2,
+        stride_t: int = 2,
+        width: int = 512,
+        depth: int = 3,
+        dilation_growth_rate: int = 3,
+        activation: str = 'relu',
+        norm: str = None,
+        kernel_size: int = 3
+    ):
+        super().__init__()
+        
+        blocks = []
+        filter_t, pad_t = stride_t * 2, stride_t // 2
+        
+        # model.0: input conv
+        blocks.append(nn.Conv1d(input_emb_width, width, kernel_size, 1, (kernel_size - 1) // 2))
+        # model.1: ReLU
+        blocks.append(nn.ReLU())
+        
+        # model.2, model.3, ...: downsample blocks
+        for i in range(down_t):
+            input_dim = width
+            block = nn.Sequential(
+                nn.Conv1d(input_dim, width, filter_t, stride_t, pad_t),
+                Resnet1D(width, depth, dilation_growth_rate, activation=activation, 
+                        norm=norm, kernel_size=kernel_size),
+            )
+            blocks.append(block)
+        
+        # model.4: output conv
+        blocks.append(nn.Conv1d(width, output_emb_width, kernel_size, 1, (kernel_size - 1) // 2))
+        
+        # MUST be named 'model' to match checkpoint keys
+        self.model = nn.Sequential(*blocks)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+
+class Decoder(nn.Module):
+    """1D CNN Decoder - matches original models/encdec.py exactly.
+    
+    Uses self.model = nn.Sequential(*blocks) to match checkpoint key structure:
+    - model.0: Conv1d (input projection)
+    - model.1: ReLU
+    - model.2: Sequential(Resnet1D, Upsample, Conv1d) for first upsample
+    - model.3: Sequential(Resnet1D, Upsample, Conv1d) for second upsample
+    - model.4: Conv1d
+    - model.5: ReLU
+    - model.6: Conv1d (output projection)
+    """
+    
+    def __init__(
+        self,
+        input_emb_width: int = 272,
+        output_emb_width: int = 512,
+        down_t: int = 2,
+        stride_t: int = 2,
+        width: int = 512,
+        depth: int = 3,
+        dilation_growth_rate: int = 3,
+        activation: str = 'relu',
+        norm: str = None,
+        kernel_size: int = 3
+    ):
+        super().__init__()
+        
+        blocks = []
+        filter_t, pad_t = stride_t * 2, stride_t // 2
+        
+        # model.0: input conv
+        blocks.append(nn.Conv1d(output_emb_width, width, kernel_size, 1, (kernel_size - 1) // 2))
+        # model.1: ReLU
+        blocks.append(nn.ReLU())
+        
+        # model.2, model.3, ...: upsample blocks
+        for i in range(down_t):
+            out_dim = width
+            block = nn.Sequential(
+                Resnet1D(width, depth, dilation_growth_rate, reverse_dilation=True,
+                        activation=activation, norm=norm, kernel_size=kernel_size),
+                nn.Upsample(scale_factor=2, mode='linear', align_corners=False),
+                nn.Conv1d(width, out_dim, 3, 1, 1)
+            )
+            blocks.append(block)
+        
+        # model.4: conv
+        blocks.append(nn.Conv1d(width, width, kernel_size, 1, (kernel_size - 1) // 2))
+        # model.5: ReLU
+        blocks.append(nn.ReLU())
+        # model.6: output conv
+        blocks.append(nn.Conv1d(width, input_emb_width, kernel_size, 1, (kernel_size - 1) // 2))
+        
+        # MUST be named 'model' to match checkpoint keys
+        self.model = nn.Sequential(*blocks)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+
+# =============================================================================
+# Vector Quantizer
+# =============================================================================
+
+class VectorQuantizerEMA(nn.Module):
+    """Single Vector Quantizer with EMA updates.
+    
+    Uses self.codebook as nn.Parameter to match checkpoint keys.
+    """
+    
+    def __init__(
+        self,
+        num_embeddings: int = 512,
+        embedding_dim: int = 8,  # per-group dimension
+        decay: float = 0.99,
+        epsilon: float = 1e-5
+    ):
+        super().__init__()
+        
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.decay = decay
+        self.epsilon = epsilon
+        
+        # MUST be named 'codebook' to match checkpoint keys
+        self.codebook = nn.Parameter(torch.randn(num_embeddings, embedding_dim))
+    
+    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            z: (B, D, T) latent features for this group
+        Returns:
+            z_q: (B, D, T) quantized features
+            indices: (B, T) codebook indices
+        """
+        B, D, T = z.shape
+        
+        # Reshape: (B, D, T) -> (B*T, D)
+        z_flat = z.permute(0, 2, 1).reshape(-1, D)
+        
+        # Compute distances to codebook
+        # d(z, e) = ||z||^2 + ||e||^2 - 2*z*e
+        distances = (
+            torch.sum(z_flat ** 2, dim=1, keepdim=True)
+            + torch.sum(self.codebook ** 2, dim=1)
+            - 2 * torch.matmul(z_flat, self.codebook.t())
+        )
+        
+        # Get nearest codebook entry
+        indices = torch.argmin(distances, dim=1)
+        
+        # Quantize
+        z_q_flat = F.embedding(indices, self.codebook)
+        
+        # Reshape back: (B*T, D) -> (B, D, T)
+        z_q = z_q_flat.reshape(B, T, D).permute(0, 2, 1)
+        
+        # Straight-through estimator
+        z_q = z + (z_q - z).detach()
+        
+        # Reshape indices: (B*T,) -> (B, T)
+        indices = indices.reshape(B, T)
+        
+        return z_q, indices
+    
+    def decode_indices(self, indices: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            indices: (B, T) codebook indices
+        Returns:
+            z_q: (B, D, T) quantized features
+        """
+        B, T = indices.shape
+        z_q_flat = F.embedding(indices.reshape(-1), self.codebook)
+        z_q = z_q_flat.reshape(B, T, -1).permute(0, 2, 1)
+        return z_q
+
+
+class MultiGroupVectorQuantizer(nn.Module):
+    """Multi-Group Vector Quantizer - splits latent into groups.
+    
+    Uses self.quantizers = nn.ModuleList to match checkpoint keys like
+    'quantizer.quantizers.0.codebook', 'quantizer.quantizers.1.codebook', etc.
+    """
+    
+    def __init__(
+        self,
+        num_groups: int = 64,
+        num_embeddings: int = 512,
+        embedding_dim: int = 512,  # total latent dim
+        decay: float = 0.99,
+        epsilon: float = 1e-5
+    ):
+        super().__init__()
+        
+        self.num_groups = num_groups
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.group_dim = embedding_dim // num_groups
+        
+        # MUST be named 'quantizers' to match checkpoint keys
+        self.quantizers = nn.ModuleList([
+            VectorQuantizerEMA(num_embeddings, self.group_dim, decay, epsilon)
+            for _ in range(num_groups)
+        ])
+    
+    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            z: (B, D, T) latent features
+        Returns:
+            z_q: (B, D, T) quantized features
+            indices: (num_groups, B, T) codebook indices per group
+        """
+        B, D, T = z.shape
+        
+        # Split into groups
+        z_groups = z.chunk(self.num_groups, dim=1)  # list of (B, group_dim, T)
+        
+        z_q_groups = []
+        indices_list = []
+        
+        for i, (z_g, quantizer) in enumerate(zip(z_groups, self.quantizers)):
+            z_q_g, idx_g = quantizer(z_g)
+            z_q_groups.append(z_q_g)
+            indices_list.append(idx_g)
+        
+        # Concatenate quantized groups
+        z_q = torch.cat(z_q_groups, dim=1)
+        
+        # Stack indices: list of (B, T) -> (num_groups, B, T)
+        indices = torch.stack(indices_list, dim=0)
+        
+        return z_q, indices
+    
+    def decode_indices(self, indices: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            indices: (num_groups, B, T) codebook indices
+        Returns:
+            z_q: (B, D, T) quantized features
+        """
+        z_q_groups = []
+        for i, quantizer in enumerate(self.quantizers):
+            z_q_g = quantizer.decode_indices(indices[i])
+            z_q_groups.append(z_q_g)
+        
+        z_q = torch.cat(z_q_groups, dim=1)
+        return z_q
+
+
+# =============================================================================
+# Main Model
+# =============================================================================
+
+class MotionVQVAE(nn.Module):
+    """Motion Vector Quantized VAE for HuggingFace.
+    
+    Architecture matches the original training code exactly to ensure
+    checkpoint compatibility.
+    """
+    
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        super().__init__()
+        
+        # Default config
+        if config is None:
+            config = {}
+        
+        self.config = config
+        
+        # Model parameters
+        # Support both naming conventions for config keys
+        self.motion_dim = config.get('motion_dim', config.get('input_dim', 272))
+        self.latent_dim = config.get('latent_dim', config.get('code_dim', 512))
+        self.num_groups = config.get('num_groups', 64)
+        self.num_codes = config.get('num_codes', config.get('nb_code', 512))
+        self.down_t = config.get('down_t', 2)
+        self.stride_t = config.get('stride_t', 2)
+        self.width = config.get('width', 512)
+        self.depth = config.get('depth', 3)
+        self.dilation_growth_rate = config.get('dilation_growth_rate', 3)
+        self.activation = config.get('activation', 'relu')
+        self.kernel_size = config.get('kernel_size', 3)
+        
+        # Normalization stats (loaded separately)
+        self.register_buffer('mean', torch.zeros(self.motion_dim))
+        self.register_buffer('std', torch.ones(self.motion_dim))
+        
+        # Build model components - names must match checkpoint
+        self.encoder = Encoder(
+            input_emb_width=self.motion_dim,
+            output_emb_width=self.latent_dim,
+            down_t=self.down_t,
+            stride_t=self.stride_t,
+            width=self.width,
+            depth=self.depth,
+            dilation_growth_rate=self.dilation_growth_rate,
+            activation=self.activation,
+            kernel_size=self.kernel_size
+        )
+        
+        self.decoder = Decoder(
+            input_emb_width=self.motion_dim,
+            output_emb_width=self.latent_dim,
+            down_t=self.down_t,
+            stride_t=self.stride_t,
+            width=self.width,
+            depth=self.depth,
+            dilation_growth_rate=self.dilation_growth_rate,
+            activation=self.activation,
+            kernel_size=self.kernel_size
+        )
+        
+        self.quantizer = MultiGroupVectorQuantizer(
+            num_groups=self.num_groups,
+            num_embeddings=self.num_codes,
+            embedding_dim=self.latent_dim
+        )
+    
+    def normalize(self, motion: torch.Tensor) -> torch.Tensor:
+        """Normalize motion data using mean and std."""
+        # motion: (B, T, D) or (B, D, T)
+        mean = self.mean.view(1, 1, -1)
+        std = self.std.view(1, 1, -1)
+        std_safe = torch.clamp(std, min=0.01)
+        
+        if motion.shape[-1] != self.motion_dim:
+            # (B, D, T) format
+            mean = mean.permute(0, 2, 1)
+            std_safe = std_safe.permute(0, 2, 1)
+        
+        normalized = (motion - mean) / std_safe
+        return torch.clamp(normalized, -20, 20)
+    
+    def denormalize(self, motion: torch.Tensor) -> torch.Tensor:
+        """Denormalize motion data using mean and std."""
+        mean = self.mean.view(1, 1, -1)
+        std = self.std.view(1, 1, -1)
+        std_safe = torch.clamp(std, min=0.01)
+        
+        if motion.shape[-1] != self.motion_dim:
+            # (B, D, T) format
+            mean = mean.permute(0, 2, 1)
+            std_safe = std_safe.permute(0, 2, 1)
+        
+        return motion * std_safe + mean
+    
+    def encode(self, motion: torch.Tensor, normalize: bool = True) -> torch.Tensor:
+        """
+        Encode motion to discrete tokens.
+        
+        Args:
+            motion: (B, T, D) motion data where D=272
+            normalize: whether to normalize input
+            
+        Returns:
+            tokens: (num_groups, B, T') discrete tokens where T' = T // 4
+        """
+        # Normalize if needed
+        if normalize:
+            motion = self.normalize(motion)
+        
+        # Convert to (B, D, T) for conv layers
+        x = motion.permute(0, 2, 1)
+        
+        # Encode
+        z = self.encoder(x)
+        
+        # Quantize
+        _, indices = self.quantizer(z)
+        
+        return indices
+    
+    def decode(self, tokens: torch.Tensor, denormalize: bool = True) -> torch.Tensor:
+        """
+        Decode discrete tokens to motion.
+        
+        Args:
+            tokens: (num_groups, B, T') discrete tokens
+            denormalize: whether to denormalize output
+            
+        Returns:
+            motion: (B, T, D) reconstructed motion
+        """
+        # Decode tokens to latent
+        z_q = self.quantizer.decode_indices(tokens)
+        
+        # Decode latent to motion
+        x_recon = self.decoder(z_q)
+        
+        # Convert to (B, T, D)
+        motion = x_recon.permute(0, 2, 1)
+        
+        # Denormalize if needed
+        if denormalize:
+            motion = self.denormalize(motion)
+        
+        return motion
+    
+    def forward(
+        self, 
+        motion: torch.Tensor, 
+        normalize: bool = True,
+        denormalize: bool = True
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Full forward pass: encode to tokens and decode back.
+        
+        Args:
+            motion: (B, T, D) motion data
+            normalize: whether to normalize input
+            denormalize: whether to denormalize output
+            
+        Returns:
+            motion_recon: (B, T, D) reconstructed motion
+            tokens: (num_groups, B, T') discrete tokens
+        """
+        # Normalize if needed
+        if normalize:
+            motion_normalized = self.normalize(motion)
+        else:
+            motion_normalized = motion
+        
+        # Convert to (B, D, T) for conv layers
+        x = motion_normalized.permute(0, 2, 1)
+        
+        # Encode
+        z = self.encoder(x)
+        
+        # Quantize
+        z_q, indices = self.quantizer(z)
+        
+        # Decode
+        x_recon = self.decoder(z_q)
+        
+        # Convert to (B, T, D)
+        motion_recon = x_recon.permute(0, 2, 1)
+        
+        # Denormalize if needed
+        if denormalize:
+            motion_recon = self.denormalize(motion_recon)
+        
+        return motion_recon, indices
+    
+    @classmethod
+    def from_pretrained(
+        cls, 
+        pretrained_path: str,
+        device: Optional[str] = None,
+        **kwargs
+    ) -> "MotionVQVAE":
+        """
+        Load pretrained model from HuggingFace Hub or local path.
+        
+        Args:
+            pretrained_path: HuggingFace repo ID (e.g., "khania/motion-vqvae") 
+                           or local directory path
+            device: Device to load model on ('cuda', 'cpu', or None for auto)
+            **kwargs: Additional arguments passed to model initialization
+            
+        Returns:
+            Loaded MotionVQVAE model
+        """
+        # Determine if path is HF repo or local
+        if os.path.isdir(pretrained_path):
+            model_dir = pretrained_path
+        elif HF_HUB_AVAILABLE:
+            model_dir = snapshot_download(repo_id=pretrained_path)
+        else:
+            raise ValueError(
+                f"Path {pretrained_path} is not a local directory and "
+                "huggingface_hub is not installed. Install with: pip install huggingface_hub"
+            )
+        
+        # Load config
+        config_path = os.path.join(model_dir, "config.json")
+        if os.path.exists(config_path):
+            with open(config_path, 'r') as f:
+                config = json.load(f)
+        else:
+            config = {}
+        
+        # Override config with kwargs
+        config.update(kwargs)
+        
+        # Create model
+        model = cls(config)
+        
+        # Load weights (includes mean/std buffers if converted with updated script)
+        weights_path = os.path.join(model_dir, "pytorch_model.bin")
+        if os.path.exists(weights_path):
+            state_dict = torch.load(weights_path, map_location='cpu')
+            
+            # Load state dict
+            missing, unexpected = model.load_state_dict(state_dict, strict=False)
+            
+            # Filter out expected missing keys (mean/std might be in separate files for older checkpoints)
+            missing_filtered = [k for k in missing if k not in ['mean', 'std']]
+            
+            if missing_filtered:
+                print(f"Warning: Missing keys in state_dict: {missing_filtered}")
+            if unexpected:
+                print(f"Warning: Unexpected keys in state_dict: {unexpected}")
+        else:
+            raise ValueError(f"Model weights not found at {weights_path}")
+        
+        # Fallback: Load mean/std from numpy files if not in state_dict
+        # This supports both old format (separate files) and new format (in weights)
+        mean_path = os.path.join(model_dir, "mean.npy")
+        std_path = os.path.join(model_dir, "std.npy")
+        
+        if 'mean' not in state_dict and os.path.exists(mean_path):
+            model.mean = torch.from_numpy(np.load(mean_path)).float()
+        if 'std' not in state_dict and os.path.exists(std_path):
+            model.std = torch.from_numpy(np.load(std_path)).float()
+        
+        # Move to device
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        model = model.to(device)
+        model.eval()
+        
+        return model
+    
+    def save_pretrained(self, save_dir: str):
+        """
+        Save model to directory in HuggingFace format.
+        
+        Args:
+            save_dir: Directory to save model to
+        """
+        os.makedirs(save_dir, exist_ok=True)
+        
+        # Save config
+        config_path = os.path.join(save_dir, "config.json")
+        with open(config_path, 'w') as f:
+            json.dump(self.config, f, indent=2)
+        
+        # Save normalization stats
+        mean_path = os.path.join(save_dir, "mean.npy")
+        std_path = os.path.join(save_dir, "std.npy")
+        np.save(mean_path, self.mean.cpu().numpy())
+        np.save(std_path, self.std.cpu().numpy())
+        
+        # Save weights
+        weights_path = os.path.join(save_dir, "pytorch_model.bin")
+        torch.save(self.state_dict(), weights_path)
+        
+        print(f"Model saved to {save_dir}")
+
+
+# =============================================================================
+# Utility Functions
+# =============================================================================
+
+def load_motion_vqvae(
+    pretrained_path: str = "khania/motion-vqvae",
+    device: Optional[str] = None
+) -> MotionVQVAE:
+    """
+    Convenience function to load MotionVQVAE.
+    
+    Args:
+        pretrained_path: HuggingFace repo ID or local path
+        device: Device to load on
+        
+    Returns:
+        Loaded model
+    """
+    return MotionVQVAE.from_pretrained(pretrained_path, device=device)
+
+
+if __name__ == "__main__":
+    # Test model creation and forward pass
+    print("Testing MotionVQVAE...")
+    
+    config = {
+        'motion_dim': 272,
+        'latent_dim': 512,
+        'num_groups': 64,
+        'num_codes': 512,
+        'down_t': 2,
+        'stride_t': 2,
+        'width': 512,
+        'depth': 3,
+        'dilation_growth_rate': 3,
+        'activation': 'relu',
+        'kernel_size': 3
+    }
+    
+    model = MotionVQVAE(config)
+    print(f"Model created with {sum(p.numel() for p in model.parameters()):,} parameters")
+    
+    # Print model architecture for debugging
+    print("\nModel state_dict keys:")
+    for k in sorted(model.state_dict().keys())[:20]:
+        print(f"  {k}")
+    print("  ...")
+    
+    # Test forward pass
+    batch_size = 2
+    seq_len = 64
+    motion = torch.randn(batch_size, seq_len, 272)
+    
+    model.eval()
+    with torch.no_grad():
+        motion_recon, tokens = model(motion, normalize=False, denormalize=False)
+    
+    print(f"\nInput shape: {motion.shape}")
+    print(f"Output shape: {motion_recon.shape}")
+    print(f"Tokens shape: {tokens.shape}")
+    print(f"MSE (random weights): {F.mse_loss(motion, motion_recon).item():.4f}")

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eaee73dd7bd25dbd49320fa242bea134efde8ad12b3a64974ea1154b8bfa2d2c
+size 293120435

std.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86c79a66805f80a5219047235536aee339de3accc4aa6de4a1857ff6ff61fc41
+size 1216