Upload 11 files

000001.npy

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

000012.npy

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

Mean.npy

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

README.md

@@ -1,3 +1,102 @@
----
-license: cc-by-nc-4.0
----
+---
+library_name: pytorch
+tags:
+- motion
+- rvq
+- vector-quantization
+- human-motion
+- safetensors
+- humanML
+- motion-reconstructor
+license: cc-by-nc-4.0
+datasets:
+- Wojtekb30/HumanML3D-500ms-FPP-descriptions-CoTs-1
+---
+
+# Motion RVQ (Move Reconstruction)
+
+This model uses Residual Vector Quantization (RVQ) to reconstruct motion represented as 263-dimensional frame vectors.
+It is a custom PyTorch model (not a Transformers `AutoModel`) and is loaded from `safetensors` with `rvq_model.py`.
+
+## Model Summary
+
+- Architecture: encoder -> 4-level RVQ -> decoder
+- Input shape: `(T, 263)` per sequence (frame-major)
+- Training window: 100 frames (with crop/pad in dataset loader)
+- Output: reconstructed motion sequence in the same 263-dim representation
+
+## Repository Files
+
+- `motion_rvq_weights.safetensors` - main published checkpoint
+- `config.json` - model configuration metadata
+- `rvq_model.py` - model architecture (`MotionRVQ_VAE`)
+- `TestRVQ.py` - inference + 3-panel visualization
+- `TrainRVQ.py` - training script
+- `rvq_humanml_dataset.py` - training dataset loader
+- `Mean.npy`, `Std.npy` - normalization statistics
+- `000001.npy`, `000012.npy` - sample motion files
+
+`motion_rvq_weights.pth` can be treated as a legacy artifact; code uses `motion_rvq_weights.safetensors`.
+
+## Install
+
+```bash
+pip install torch safetensors numpy matplotlib
+```
+
+## Inference
+
+Run the provided visualization script:
+
+```bash
+python TestRVQ.py
+```
+
+By default, `TestRVQ.py` uses `000001.npy`. You can change `FILE_TO_TEST` in `TestRVQ.py` to another sequence.
+
+Minimal loading example:
+
+```python
+from pathlib import Path
+import torch
+from safetensors.torch import load_file
+from rvq_model import MotionRVQ_VAE
+
+base = Path(".")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = MotionRVQ_VAE().to(device)
+state_dict = load_file(str(base / "motion_rvq_weights.safetensors"), device=str(device))
+model.load_state_dict(state_dict)
+model.eval()
+```
+
+## Training From Scratch
+
+Expected layout:
+
+```text
+rvq/
+  TrainRVQ.py
+  rvq_model.py
+  rvq_humanml_dataset.py
+  Mean.npy
+  Std.npy
+  new_joint_vecs/
+    *.npy
+```
+
+Run training:
+
+```bash
+python TrainRVQ.py
+```
+
+Output checkpoint:
+
+- `motion_rvq_weights.safetensors`
+
+## Limitations
+
+- This model reconstructs motion vectors; it is not a text-to-motion generator.
+- Input format must match the same 263-dim representation and normalization scheme used during training.

Std.npy

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

TestRVQ.py

@@ -0,0 +1,144 @@
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+import torch.nn.functional as F
+from pathlib import Path
+from safetensors.torch import load_file
+
+# Import model architecture
+from rvq_model import MotionRVQ_VAE
+
+# ==========================================
+# 1. Configuration
+# ==========================================
+BASE_DIR = Path(__file__).resolve().parent
+FILE_TO_TEST = BASE_DIR / "000001.npy"
+WEIGHTS_PATH = BASE_DIR / "motion_rvq_weights.safetensors"
+
+# ==========================================
+# 2. Model Initialization
+# ==========================================
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = MotionRVQ_VAE().to(device)
+
+try:
+    state_dict = load_file(str(WEIGHTS_PATH), device=str(device))
+    model.load_state_dict(state_dict)
+    print(f"Successfully loaded weights from {WEIGHTS_PATH}!")
+except FileNotFoundError:
+    print(f"ERROR: Could not find file {WEIGHTS_PATH}.")
+    exit()
+
+model.eval()
+
+# ==========================================
+# 3. Forward Pass (With Padding and Normalization)
+# ==========================================
+original_data = np.load(FILE_TO_TEST)
+T_orig = original_data.shape[0]
+
+# Load normalization vectors
+mean = np.load(BASE_DIR / "Mean.npy")
+std = np.load(BASE_DIR / "Std.npy")
+
+# Padding for stride=4 compression
+pad_len = (4 - (T_orig % 4)) % 4
+if pad_len > 0:
+    last_frame = original_data[-1:]
+    padded_data = np.concatenate([original_data, np.repeat(last_frame, pad_len, axis=0)], axis=0)
+else:
+    padded_data = original_data
+
+padded_data = (padded_data - mean) / std
+input_tensor = torch.from_numpy(padded_data).float().unsqueeze(0).permute(0, 2, 1).to(device)
+
+with torch.no_grad():
+    # 1. Get tokens from all levels
+    z = model.encoder(input_tensor)
+    _, tokens, _ = model.rvq(z)
+    
+    # 2. Function for "partial" decoding
+    def decode_from_levels(num_levels):
+        z_q_partial = 0
+        for i in range(num_levels):
+            # Get indices only for level "i"
+            indices = tokens[:, i, :] 
+            quantizer = model.rvq.quantizers[i]
+            
+            # Convert token id (e.g. 145) into its 1024-codebook vector
+            level_z_q = F.embedding(indices, quantizer.embedding)
+            level_z_q = level_z_q.permute(0, 2, 1)  # Shape expected by decoder
+            
+            # Add residual vector to the running latent
+            z_q_partial = z_q_partial + level_z_q
+            
+        return model.decoder(z_q_partial)
+
+    # Generate motion using only level 1 and all 4 levels
+    recon_tensor_1_lvl = decode_from_levels(1)
+    recon_tensor_4_lvl = decode_from_levels(4)
+
+# Back to NumPy
+recon_1_lvl = recon_tensor_1_lvl.squeeze(0).permute(1, 0).cpu().numpy()[:T_orig, :]
+recon_4_lvl = recon_tensor_4_lvl.squeeze(0).permute(1, 0).cpu().numpy()[:T_orig, :]
+
+# De-normalization
+recon_1_lvl = (recon_1_lvl * std) + mean
+recon_4_lvl = (recon_4_lvl * std) + mean
+
+# ==========================================
+# 4. Skeleton extraction utility
+# ==========================================
+def get_3d_joints(data_263):
+    num_frames = data_263.shape[0]
+    joints = np.zeros((num_frames, 22, 3))
+    for i in range(num_frames):
+        root_y = data_263[i, 3]
+        joints[i, 0] = [0, root_y, 0] 
+        local_positions = data_263[i, 4:67].reshape(21, 3)
+        joints[i, 1:] = local_positions + [0, root_y, 0]
+    return joints
+
+joints_orig = get_3d_joints(original_data)
+joints_1_lvl = get_3d_joints(recon_1_lvl)
+joints_4_lvl = get_3d_joints(recon_4_lvl)
+
+# ==========================================
+# 5. Three-panel visualization
+# ==========================================
+kinematic_tree = [
+    [0, 1, 4, 7, 10], [0, 2, 5, 8, 11], [0, 3, 6, 9, 12, 15], 
+    [9, 13, 16, 18, 20], [9, 14, 17, 19, 21]
+]
+
+fig = plt.figure(figsize=(15, 5))
+ax1 = fig.add_subplot(131, projection='3d')
+ax2 = fig.add_subplot(132, projection='3d')
+ax3 = fig.add_subplot(133, projection='3d')
+
+def update(frame):
+    for ax in [ax1, ax2, ax3]:
+        ax.clear()
+        ax.set_xlim(-1, 1); ax.set_ylim(-1, 1); ax.set_zlim(0, 2)
+        ax.view_init(elev=10., azim=-90)
+        ax.axis('off')
+        
+    ax1.set_title(f"ORIGINAL\n(Frame {frame})")
+    ax2.set_title("RECONSTRUCTION: 1 LEVEL\n(Coarse tokens only)")
+    ax3.set_title("RECONSTRUCTION: 4 LEVELS\n(Full RVQ detail)")
+    
+    for chain in kinematic_tree:
+        # Original (Blue)
+        ax1.plot(joints_orig[frame, chain, 0], joints_orig[frame, chain, 2], joints_orig[frame, chain, 1], 
+                 linewidth=3, marker='o', markersize=4, color='blue')
+        # 1 Level (Orange)
+        ax2.plot(joints_1_lvl[frame, chain, 0], joints_1_lvl[frame, chain, 2], joints_1_lvl[frame, chain, 1], 
+                 linewidth=3, marker='o', markersize=4, color='orange')
+        # 4 Levels (Red)
+        ax3.plot(joints_4_lvl[frame, chain, 0], joints_4_lvl[frame, chain, 2], joints_4_lvl[frame, chain, 1], 
+                 linewidth=3, marker='o', markersize=4, color='red')
+
+ani = animation.FuncAnimation(fig, update, frames=T_orig, interval=50, repeat=True)
+plt.tight_layout()
+plt.show()

TrainRVQ.py

@@ -0,0 +1,68 @@
+import multiprocessing
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+from safetensors.torch import save_file
+
+from rvq_model import MotionRVQ_VAE
+
+
+if __name__ == "__main__":
+    multiprocessing.freeze_support()
+
+    from rvq_humanml_dataset import DataLoader, HumanML3DDataset
+
+    base_dir = Path(__file__).resolve().parent
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Training on: {device}")
+
+    model = MotionRVQ_VAE().to(device)
+
+    dataset = HumanML3DDataset(data_dir=str(base_dir / "new_joint_vecs"), window_size=100)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=32,
+        shuffle=True,
+        num_workers=4,
+        drop_last=True,
+    )
+
+    optimizer = optim.AdamW(model.parameters(), lr=2e-4, weight_decay=1e-4)
+    num_epochs = 5
+
+    model.train()
+    for epoch in range(num_epochs):
+        epoch_loss = 0.0
+
+        for batch_idx, batch in enumerate(dataloader):
+            batch = batch.to(device)
+            optimizer.zero_grad()
+
+            reconstructed, _, commit_loss = model(batch)
+            pos_loss = F.mse_loss(reconstructed, batch)
+
+            vel_orig = batch[:, :, 1:] - batch[:, :, :-1]
+            vel_recon = reconstructed[:, :, 1:] - reconstructed[:, :, :-1]
+            vel_loss = F.mse_loss(vel_recon, vel_orig)
+
+            reconstruction_loss = pos_loss + (1.5 * vel_loss)
+            loss = reconstruction_loss + commit_loss
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+
+            if batch_idx % 50 == 0:
+                print(
+                    f"Epoch [{epoch + 1}/{num_epochs}] Batch [{batch_idx}/{len(dataloader)}] "
+                    f"MSE: {reconstruction_loss.item():.4f} | Commit: {commit_loss.item():.4f}"
+                )
+
+        print(f"--- End of epoch {epoch + 1} | Avg loss: {epoch_loss / len(dataloader):.4f} ---")
+
+    weights_path = base_dir / "motion_rvq_weights.safetensors"
+    state_dict = {k: v.detach().cpu() for k, v in model.state_dict().items()}
+    save_file(state_dict, str(weights_path))
+    print(f"Training complete and model saved to: {weights_path}")

config.json

@@ -0,0 +1,23 @@
+{
+  "architectures": [
+    "MotionRVQ_VAE"
+  ],
+  "model_type": "motion_rvq_vae",
+  "library_name": "pytorch",
+  "torch_dtype": "float32",
+  "in_channels": 263,
+  "out_channels": 263,
+  "latent_dim": 512,
+  "encoder_hidden_channels": 512,
+  "decoder_hidden_channels": 512,
+  "rvq_num_levels": 4,
+  "rvq_num_embeddings": 1024,
+  "rvq_embedding_dim": 512,
+  "rvq_commitment_cost": 0.25,
+  "rvq_decay": 0.99,
+  "rvq_epsilon": 1e-5,
+  "downsample_stride": 4,
+  "window_size": 100,
+  "input_representation": "263-dim frame vectors",
+  "task": "motion_reconstruction"
+}

motion_rvq_weights.safetensors

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

rvq_humanml_dataset.py

@@ -0,0 +1,42 @@
+import os
+import glob
+from pathlib import Path
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+
+class HumanML3DDataset(Dataset):
+    def __init__(self, data_dir='./new_joint_vecs', window_size=100):
+        self.data_dir = str(data_dir)
+        self.window_size = window_size
+        self.file_paths = glob.glob(os.path.join(self.data_dir, '*.npy'))
+
+        data_dir_path = Path(self.data_dir).resolve()
+        base_dir = data_dir_path.parent
+        self.mean = torch.from_numpy(np.load(base_dir / 'Mean.npy')).float()
+        self.std = torch.from_numpy(np.load(base_dir / 'Std.npy')).float()
+
+        print(f"Dataset initialized: found {len(self.file_paths)} files.")
+
+    def __len__(self):
+        return len(self.file_paths)
+
+    def __getitem__(self, idx):
+        data = np.load(self.file_paths[idx])
+        tensor_data = torch.from_numpy(data).float()
+
+        tensor_data = (tensor_data - self.mean) / self.std
+
+        t_len, _ = tensor_data.shape
+        if t_len > self.window_size:
+            start = torch.randint(0, t_len - self.window_size + 1, (1,)).item()
+            tensor_data = tensor_data[start : start + self.window_size, :]
+        elif t_len < self.window_size:
+            padding_size = self.window_size - t_len
+            last_pose = tensor_data[-1].unsqueeze(0)
+            padding = last_pose.repeat(padding_size, 1)
+            tensor_data = torch.cat([tensor_data, padding], dim=0)
+
+        tensor_data = tensor_data.permute(1, 0)
+        return tensor_data

rvq_model.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class EMAVectorQuantizer(nn.Module):
+    def __init__(
+        self,
+        num_embeddings=512,
+        embedding_dim=256,
+        commitment_cost=0.25,
+        decay=0.99,
+        epsilon=1e-5,
+    ):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.commitment_cost = commitment_cost
+        self.decay = decay
+        self.epsilon = epsilon
+
+        embed = torch.randn(num_embeddings, embedding_dim)
+        self.register_buffer("embedding", embed)
+        self.register_buffer("cluster_size", torch.zeros(num_embeddings))
+        self.register_buffer("ema_w", embed.clone())
+
+    def forward(self, z):
+        z = z.permute(0, 2, 1).contiguous()
+        z_flattened = z.view(-1, self.embedding_dim)
+
+        distances = (
+            torch.sum(z_flattened**2, dim=1, keepdim=True)
+            + torch.sum(self.embedding**2, dim=1)
+            - 2 * torch.matmul(z_flattened, self.embedding.t())
+        )
+
+        min_encoding_indices = torch.argmin(distances, dim=1)
+        z_q = F.embedding(min_encoding_indices, self.embedding)
+
+        if self.training:
+            encodings = F.one_hot(min_encoding_indices, self.num_embeddings).float()
+            self.cluster_size.data.mul_(self.decay).add_(encodings.sum(0), alpha=1 - self.decay)
+
+            n = self.cluster_size.sum()
+            cluster_size = (self.cluster_size + self.epsilon) / (
+                n + self.num_embeddings * self.epsilon
+            ) * n
+
+            dw = torch.matmul(encodings.t(), z_flattened)
+            self.ema_w.data.mul_(self.decay).add_(dw, alpha=1 - self.decay)
+            self.embedding.data.copy_(self.ema_w / cluster_size.unsqueeze(1))
+
+        loss = self.commitment_cost * F.mse_loss(z_q.detach(), z_flattened)
+        z_q = z_flattened + (z_q - z_flattened).detach()
+        z_q = z_q.view(z.shape).permute(0, 2, 1).contiguous()
+
+        return z_q, min_encoding_indices.view(z.shape[0], z.shape[1]), loss
+
+
+class RVQ(nn.Module):
+    def __init__(self, num_levels=3, num_embeddings=512, embedding_dim=256):
+        super().__init__()
+        self.num_levels = num_levels
+        self.quantizers = nn.ModuleList(
+            [EMAVectorQuantizer(num_embeddings, embedding_dim) for _ in range(num_levels)]
+        )
+
+    def forward(self, z):
+        quantized_out = 0
+        residual = z
+        all_indices = []
+        total_loss = 0
+
+        for quantizer in self.quantizers:
+            z_q, indices, loss = quantizer(residual)
+            quantized_out = quantized_out + z_q
+            residual = residual - z_q
+            all_indices.append(indices)
+            total_loss += loss
+
+        return quantized_out, torch.stack(all_indices, dim=1), total_loss
+
+
+class ResBlock1D(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv1d(channels, channels, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv1d(channels, channels, kernel_size=3, padding=1),
+        )
+
+    def forward(self, x):
+        return x + self.net(x)
+
+
+class MotionEncoder(nn.Module):
+    def __init__(self, in_channels=263, latent_dim=512):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv1d(in_channels, 512, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2, inplace=True),
+            ResBlock1D(512),
+            ResBlock1D(512),
+            ResBlock1D(512),
+            nn.Conv1d(512, latent_dim, kernel_size=8, stride=4, padding=2),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class MotionDecoder(nn.Module):
+    def __init__(self, latent_dim=512, out_channels=263):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.ConvTranspose1d(latent_dim, 512, kernel_size=8, stride=4, padding=2),
+            nn.LeakyReLU(0.2, inplace=True),
+            ResBlock1D(512),
+            ResBlock1D(512),
+            ResBlock1D(512),
+            nn.Conv1d(512, out_channels, kernel_size=3, padding=1),
+        )
+
+    def forward(self, z_q):
+        return self.net(z_q)
+
+
+class MotionRVQ_VAE(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = MotionEncoder(in_channels=263, latent_dim=512)
+        self.rvq = RVQ(num_levels=4, num_embeddings=1024, embedding_dim=512)
+        self.decoder = MotionDecoder(latent_dim=512, out_channels=263)
+
+    def forward(self, x):
+        z = self.encoder(x)
+        z_q, token_indices, commitment_loss = self.rvq(z)
+        x_recon = self.decoder(z_q)
+        return x_recon, token_indices, commitment_loss