Upload folder using huggingface_hub

README.md

@@ -1,13 +1,36 @@
----
-title: Irg Motionllm
-emoji: 🦀
-colorFrom: purple
-colorTo: blue
-sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: IRG-MotionLLM Text-to-Motion
+emoji: 🕺
+colorFrom: pink
+colorTo: indigo
+sdk: gradio
+app_file: app.py
+pinned: false
+short_description: Text-to-3D human motion generation
+python_version: "3.10"
+startup_duration_timeout: 40m
+---
+
+# IRG-MotionLLM — Text-to-3D-Motion
+
+Interactive demo for
+[**IRG-MotionLLM: Interleaving Motion Generation, Assessment and Refinement
+for Text-to-Motion Generation**](https://huggingface.co/papers/2512.10730).
+
+Enter a natural-language description of a human motion and the model generates
+a 3D skeletal animation (HumanML3D, 22 joints).
+
+## How it works
+
+1. A Gemma-2-2B LLM (fine-tuned by the authors) reads the prompt and, in an
+   *interleaved* chain-of-thought, generates → self-assesses → refines discrete
+   motion tokens until it is satisfied.
+2. The final motion tokens are decoded by a VQ-VAE into 263-dimensional
+   HumanML3D motion features.
+3. The features are converted to 3D joint positions (`recover_from_ric`) and
+   rendered to a skeleton video.
+
+Model weights: [`Lymann/IRG-MotionLLM-HumanML3D`](https://huggingface.co/Lymann/IRG-MotionLLM-HumanML3D)
+· Code: [HumanMLLM/IRG-MotionLLM](https://github.com/HumanMLLM/IRG-MotionLLM)
+
+Models are released under CC-BY-NC-SA 4.0 for research use only.

app.py

@@ -0,0 +1,391 @@
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+os.environ.setdefault("MPLBACKEND", "Agg")
+
+import re
+import types
+import tempfile
+
+import spaces
+import torch
+import numpy as np
+import gradio as gr
+import matplotlib
+
+matplotlib.use("Agg")
+
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from motionvq.vqvae import HumanVQVAE
+from motionvq.motion_utils import recover_from_ric, plot_3d_motion
+from motionvq.paramUtil import t2m_kinematic_chain
+
+# ----------------------------------------------------------------------------
+# Config -- values verified against the IRG-MotionLLM / Motion-Agent source and
+# the actual Stage-3 checkpoint tensor shapes:
+#   * base LLM  : google/gemma-2-2b-it (bf16)
+#   * added toks : <Motion>, </Motion>, and 512 motion codes  (nb_code = 512)
+#   * VQVAE     : HumanML3D 263-dim, codebook 512x512, down_t=2, stride_t=2
+#   * Stage-3 checkpoint bundles BOTH the merged full LLM and the VQVAE weights
+# ----------------------------------------------------------------------------
+BASE_LLM = "google/gemma-2-2b-it"
+CKPT_REPO = "Lymann/IRG-MotionLLM-HumanML3D"
+CKPT_FILE = "Stage-3/irg_motionllm_unified_rl_stage3.bin"
+
+NB_CODE = 512
+JOINTS_NUM = 22
+FPS = 20
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+
+class VQArgs:
+    """Minimal args namespace expected by HumanVQVAE / VQVAE_251."""
+
+    dataname = "t2m"
+    quantizer = "ema_reset"
+    mu = 0.99
+
+
+VQ_ARGS = VQArgs()
+
+# System prompt for the IRG (interleaved generation-assessment-refinement) task,
+# copied verbatim from the repo's TextProcessor (unified_mogen_cot_v3).
+IRG_SYSTEM_PROMPT = (
+    "You are an assistant who helps users understand or generate 3D human "
+    "motion representations."
+)
+IRG_USER_TEMPLATE = (
+    "Given a text outlining a human motion objective, employ a step-by-step "
+    "thought process to realize the motion: (1) analyze the text, providing a "
+    "clear explanation of the reasoning to identify essential elements; "
+    "(2) conduct several rounds of motion generation and self-assessment until "
+    "the motion is satisfactory. Wrap all responses in <think> and </think> "
+    "tags, and formulate a plan before each step.\nGoal Text: {caption}"
+)
+
+# ----------------------------------------------------------------------------
+# Load model + VQVAE at module scope (ZeroGPU: .to("cuda") is intercepted)
+# ----------------------------------------------------------------------------
+print("Loading tokenizer + base LLM ...")
+tokenizer = AutoTokenizer.from_pretrained(BASE_LLM, token=HF_TOKEN)
+
+llm = AutoModelForCausalLM.from_pretrained(
+    BASE_LLM,
+    torch_dtype=torch.bfloat16,
+    attn_implementation="sdpa",
+    token=HF_TOKEN,
+)
+
+# Recreate the exact tokenizer vocabulary the model was trained with:
+#   base vocab + <Motion> + </Motion> + <Motion_0> .. <Motion_{NB_CODE-1}>
+NB_TEXT_TOKENS = len(tokenizer)
+tokenizer.add_tokens(["<Motion>", "</Motion>"])
+for i in range(NB_CODE):
+    tokenizer.add_tokens([f"<Motion_{i}>"])
+llm.resize_token_embeddings(len(tokenizer))
+
+print("Loading Stage-3 checkpoint ...")
+ckpt_path = hf_hub_download(CKPT_REPO, CKPT_FILE, token=HF_TOKEN)
+state = torch.load(ckpt_path, map_location="cpu")
+if isinstance(state, dict) and "state_dict" in state:
+    state = state["state_dict"]
+
+# The Stage-3 checkpoint is a full fine-tune (LoRA already merged) and stores
+# both the LLM weights (prefixed "llm.") and the VQVAE weights (prefixed "net.").
+llm_sd = {}
+net_sd = {}
+for k, v in state.items():
+    if k.startswith("llm."):
+        llm_sd[k[len("llm."):]] = v
+    elif k.startswith("net."):
+        net_sd[k[len("net."):]] = v
+
+missing, unexpected = llm.load_state_dict(llm_sd, strict=False)
+print(f"LLM load -> missing={len(missing)} unexpected={len(unexpected)}")
+llm = llm.to(torch.bfloat16).eval()
+
+print("Building + loading VQVAE ...")
+net = HumanVQVAE(
+    VQ_ARGS,
+    nb_code=NB_CODE,
+    code_dim=512,
+    output_emb_width=512,
+    down_t=2,
+    stride_t=2,
+    width=512,
+    depth=3,
+    dilation_growth_rate=3,
+    activation="relu",
+    norm=None,
+)
+vq_missing, vq_unexpected = net.load_state_dict(net_sd, strict=False)
+print(f"VQVAE load -> missing={len(vq_missing)} unexpected={len(vq_unexpected)}")
+net = net.eval().float()
+for p in net.parameters():
+    p.requires_grad = False
+
+# HumanML3D normalization statistics (263-dim), extracted from the same
+# T2M-GPT extractor bundle the authors use.
+MEAN = np.load(os.path.join(os.path.dirname(__file__), "mean.npy"))
+STD = np.load(os.path.join(os.path.dirname(__file__), "std.npy"))
+
+llm.to("cuda")
+net.to("cuda")
+
+EOS_ID = tokenizer.eos_token_id
+MOTION_ID_START = len(tokenizer) - (NB_CODE + 2)  # first <Motion_i> token id
+
+
+# ----------------------------------------------------------------------------
+# Motion-token extraction helpers (ported from mllm_single_lora.py)
+# ----------------------------------------------------------------------------
+def _find_seq(tokens_list, target):
+    for i in range(len(tokens_list) - len(target) + 1):
+        if tokens_list[i:i + len(target)] == target:
+            return i
+    return -1
+
+
+def _extract_answer_scores(scores, ids_list):
+    """Return the score rows and matching ids inside <answer>..</answer>,
+    else fall back to the last <Motion>..</Motion> span, else all."""
+    ans_start = tokenizer.encode("<answer>", add_special_tokens=False)
+    ans_end = tokenizer.encode("</answer>", add_special_tokens=False)
+    s = _find_seq(ids_list, ans_start)
+    if s != -1:
+        s2 = s + len(ans_start)
+        e = _find_seq(ids_list[s2:], ans_end)
+        if e != -1:
+            e += s2
+            return scores[s2:e], ids_list[s2:e]
+    return scores, ids_list
+
+
+def _tokens_from_scores_and_ids(scores, ids_list):
+    """Prefer already-decoded motion token ids that fall in the motion range;
+    otherwise argmax over the motion-logit slice. Mirrors the repo logic."""
+    motion_logits = scores[:, -(NB_CODE + 2):]
+    argmax_ids = torch.argmax(motion_logits, dim=-1)
+    out = []
+    for i, tid in enumerate(ids_list):
+        if MOTION_ID_START <= tid < len(tokenizer):
+            out.append(tid - MOTION_ID_START)
+        else:
+            out.append(int(argmax_ids[i].item()))
+    motion_tokens = torch.tensor(out, dtype=torch.long)
+    if 1 in motion_tokens.tolist():
+        motion_tokens = motion_tokens[:motion_tokens.tolist().index(1)]
+    if 0 in motion_tokens.tolist():
+        motion_tokens = motion_tokens[motion_tokens.tolist().index(0) + 1:]
+    motion_tokens = torch.clamp(motion_tokens - 2, min=0)
+    return motion_tokens
+
+
+def _extract_last_motion_span_ids(ids_list):
+    st_id = tokenizer.encode("<Motion>", add_special_tokens=False)[0]
+    ed_id = tokenizer.encode("</Motion>", add_special_tokens=False)[0]
+    starts = [i for i, t in enumerate(ids_list) if t == st_id]
+    ends = [i for i, t in enumerate(ids_list) if t == ed_id]
+    if not starts or not ends:
+        return None
+    # last valid <Motion> ... </Motion> pair
+    for st in reversed(starts):
+        later_ends = [e for e in ends if e > st]
+        if later_ends:
+            return st, later_ends[0]
+    return None
+
+
+def build_prompt(caption):
+    caption = (caption or "").strip()
+    user_prompt = IRG_USER_TEMPLATE.format(caption=caption)
+    text = IRG_SYSTEM_PROMPT + "\n\n" + "User: " + user_prompt + "\n\n" + " Response:"
+    return text
+
+
+def decode_motion_to_joints(motion_tokens):
+    """VQVAE token ids -> 263-dim features -> denormalize -> 22-joint xyz."""
+    motion_tokens = motion_tokens.to("cuda").long()
+    feats = net.forward_decoder(motion_tokens)          # (1, T, 263)
+    feats = feats.detach().cpu().numpy()[0]
+    feats = MEAN + feats * STD                          # denormalize
+    joints = recover_from_ric(torch.from_numpy(feats).float(), JOINTS_NUM)
+    return joints.numpy()                               # (T, 22, 3)
+
+
+def render_video(joints, title):
+    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    plot_3d_motion(out_path, t2m_kinematic_chain, joints, title=title, fps=FPS, radius=4)
+    return out_path
+
+
+# ----------------------------------------------------------------------------
+# Inference
+# ----------------------------------------------------------------------------
+def _duration(caption, reasoning, max_new_tokens, *a, **k):
+    base = 55
+    if reasoning:
+        base = 110
+    return min(160, base + int(int(max_new_tokens) / 20))
+
+
+@spaces.GPU(duration=_duration)
+def generate(caption, reasoning=True, max_new_tokens=1300, seed=0,
+             progress=gr.Progress(track_tqdm=True)):
+    if not caption or not caption.strip():
+        raise gr.Error("Please enter a motion description.")
+
+    if seed and int(seed) > 0:
+        torch.manual_seed(int(seed))
+
+    do_sample = bool(seed and int(seed) > 0)
+    max_new_tokens = int(max_new_tokens)
+
+    prompt = build_prompt(caption)
+    enc = tokenizer(prompt, return_tensors="pt").to("cuda")
+    input_len = enc.input_ids.shape[1]
+
+    if reasoning:
+        max_len = input_len + max_new_tokens
+    else:
+        # skip the interleaved reasoning; ask for a compact motion answer
+        max_len = input_len + 260
+
+    gen_kwargs = dict(
+        max_length=max_len,
+        do_sample=do_sample,
+        return_dict_in_generate=True,
+        output_scores=True,
+        use_cache=True,
+    )
+    if do_sample:
+        gen_kwargs["temperature"] = 1.0
+
+    with torch.inference_mode():
+        outputs = llm.generate(enc.input_ids, attention_mask=enc.attention_mask, **gen_kwargs)
+
+    gen_ids = outputs.sequences[0, input_len:]
+    # scores are per generated step (tuple len == n_new); stack into (n_new, vocab)
+    scores = torch.stack(outputs.scores)[:, 0, :]
+
+    # truncate at EOS
+    gen_ids_list = gen_ids.tolist()
+    if EOS_ID in gen_ids_list:
+        cut = gen_ids_list.index(EOS_ID) + 1
+        gen_ids_list = gen_ids_list[:cut]
+        scores = scores[:cut]
+
+    full_text = tokenizer.decode(gen_ids_list, skip_special_tokens=False)
+
+    # Prefer <answer>..</answer>; fall back to last <Motion>..</Motion> span.
+    ans_scores, ans_ids = _extract_answer_scores(scores, gen_ids_list)
+    if ans_ids is gen_ids_list or len(ans_ids) == len(gen_ids_list):
+        span = _extract_last_motion_span_ids(gen_ids_list)
+        if span is not None:
+            st, ed = span
+            ans_scores = scores[st + 1:ed]
+            ans_ids = gen_ids_list[st + 1:ed]
+
+    motion_tokens = _tokens_from_scores_and_ids(ans_scores, ans_ids)
+
+    if motion_tokens.numel() == 0:
+        raise gr.Error(
+            "The model did not produce a valid motion for this prompt. "
+            "Try rephrasing, or toggle the reasoning option."
+        )
+
+    joints = decode_motion_to_joints(motion_tokens)
+    title = caption.strip()
+    if len(title) > 60:
+        title = title[:57] + "..."
+    video = render_video(joints, title)
+
+    # Build a readable reasoning trace (strip the internal [plan]/[tag] markup).
+    trace = full_text
+    trace = trace.replace("<eos>", "").replace("<pad>", "")
+    n_rounds = len(re.findall(r"\[generate\]", trace))
+    info = (
+        f"Frames: {joints.shape[0]}  |  Motion tokens: {motion_tokens.numel()}"
+        f"  |  Internal generate/refine rounds: {max(n_rounds, 1)}"
+    )
+    return video, info, trace.strip()
+
+
+# ----------------------------------------------------------------------------
+# UI
+# ----------------------------------------------------------------------------
+DESCRIPTION = """
+# 🕺 IRG-MotionLLM — Text-to-3D-Motion
+
+Generate 3D human motion from a text description with
+**[IRG-MotionLLM](https://huggingface.co/papers/2512.10730)** — an LLM
+(Gemma-2-2B) that *interleaves* motion generation, self-assessment and
+refinement, then decodes discrete motion tokens through a VQ-VAE into a
+HumanML3D skeleton animation.
+"""
+
+EXAMPLES = [
+    ["a person walks forward, then turns around and walks back.", True, 1300, 0],
+    ["a man is doing cartwheels.", True, 1300, 0],
+    ["a person jumps up high with both hands raised.", True, 1300, 0],
+    ["someone sits down on a chair and crosses their legs.", True, 1300, 0],
+    ["a person raises their right hand and waves.", True, 1300, 0],
+]
+
+with gr.Blocks(theme=gr.themes.Citrus()) as demo:
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            caption = gr.Textbox(
+                label="Motion description",
+                placeholder="e.g. a person walks forward, then turns around and walks back.",
+                lines=2,
+            )
+            run_btn = gr.Button("Generate motion", variant="primary")
+            with gr.Accordion("Advanced options", open=False):
+                reasoning = gr.Checkbox(
+                    value=True,
+                    label="Interleaved reasoning (generate → assess → refine)",
+                    info="Uses the full IRG chain-of-thought. Turn off for a faster, direct generation.",
+                )
+                max_new_tokens = gr.Slider(
+                    400, 2000, value=1300, step=50,
+                    label="Max new tokens (reasoning budget)",
+                )
+                seed = gr.Slider(
+                    0, 100000, value=0, step=1,
+                    label="Seed (0 = greedy / deterministic, >0 = sampling)",
+                )
+        with gr.Column(scale=1):
+            video_out = gr.Video(label="Generated motion", autoplay=True)
+            info_out = gr.Textbox(label="Summary", lines=2)
+
+    with gr.Accordion("Model reasoning trace", open=False):
+        trace_out = gr.Textbox(label="Interleaved generation / assessment / refinement", lines=10)
+
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=[caption, reasoning, max_new_tokens, seed],
+        outputs=[video_out, info_out, trace_out],
+        fn=generate,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+    run_btn.click(
+        fn=generate,
+        inputs=[caption, reasoning, max_new_tokens, seed],
+        outputs=[video_out, info_out, trace_out],
+    )
+    caption.submit(
+        fn=generate,
+        inputs=[caption, reasoning, max_new_tokens, seed],
+        outputs=[video_out, info_out, trace_out],
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch()

mean.npy

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

motionvq/encdec.py

@@ -0,0 +1,67 @@
+import torch.nn as nn
+from .resnet import Resnet1D
+
+class Encoder(nn.Module):
+    def __init__(self,
+                 input_emb_width = 3,
+                 output_emb_width = 512,
+                 down_t = 3,
+                 stride_t = 2,
+                 width = 512,
+                 depth = 3,
+                 dilation_growth_rate = 3,
+                 activation='relu',
+                 norm=None):
+        super().__init__()
+        
+        blocks = []
+        filter_t, pad_t = stride_t * 2, stride_t // 2
+        blocks.append(nn.Conv1d(input_emb_width, width, 3, 1, 1))
+        blocks.append(nn.ReLU())
+        
+        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),
+            )
+            blocks.append(block)
+        blocks.append(nn.Conv1d(width, output_emb_width, 3, 1, 1))
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.model(x)
+
+class Decoder(nn.Module):
+    def __init__(self,
+                 input_emb_width = 3,
+                 output_emb_width = 512,
+                 down_t = 3,
+                 stride_t = 2,
+                 width = 512,
+                 depth = 3,
+                 dilation_growth_rate = 3, 
+                 activation='relu',
+                 norm=None):
+        super().__init__()
+        blocks = []
+        
+        filter_t, pad_t = stride_t * 2, stride_t // 2
+        blocks.append(nn.Conv1d(output_emb_width, width, 3, 1, 1))
+        blocks.append(nn.ReLU())
+        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),
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                nn.Conv1d(width, out_dim, 3, 1, 1)
+            )
+            blocks.append(block)
+        blocks.append(nn.Conv1d(width, width, 3, 1, 1))
+        blocks.append(nn.ReLU())
+        blocks.append(nn.Conv1d(width, input_emb_width, 3, 1, 1))
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.model(x)
+    

motionvq/motion_utils.py

@@ -0,0 +1,277 @@
+import torch
+import numpy as np
+import shutil as _shutil
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+_ffmpeg = _shutil.which('ffmpeg')
+if _ffmpeg:
+    plt.rcParams['animation.ffmpeg_path'] = _ffmpeg
+from mpl_toolkits.mplot3d import Axes3D
+from matplotlib.animation import FuncAnimation, PillowWriter
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+import mpl_toolkits.mplot3d.axes3d as p3
+
+def qrot(q, v):
+    """
+    Rotate vector(s) v about the rotation described by quaternion(s) q.
+    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
+    where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 3).
+    """
+    assert q.shape[-1] == 4
+    assert v.shape[-1] == 3
+    assert q.shape[:-1] == v.shape[:-1]
+
+    original_shape = list(v.shape)
+    # print(q.shape)
+    q = q.contiguous().view(-1, 4)
+    v = v.contiguous().view(-1, 3)
+
+    qvec = q[:, 1:]
+    uv = torch.cross(qvec, v, dim=1)
+    uuv = torch.cross(qvec, uv, dim=1)
+    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)
+
+def qinv(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    mask = torch.ones_like(q)
+    mask[..., 1:] = -mask[..., 1:]
+    return q * mask
+
+def recover_root_rot_pos(data):
+    rot_vel = data[..., 0]
+    r_rot_ang = torch.zeros_like(rot_vel).to(data.device)
+    '''Get Y-axis rotation from rotation velocity'''
+    r_rot_ang[..., 1:] = rot_vel[..., :-1]
+    r_rot_ang = torch.cumsum(r_rot_ang, dim=-1)
+
+    r_rot_quat = torch.zeros(data.shape[:-1] + (4,)).to(data.device)
+    r_rot_quat[..., 0] = torch.cos(r_rot_ang)
+    r_rot_quat[..., 2] = torch.sin(r_rot_ang)
+
+    r_pos = torch.zeros(data.shape[:-1] + (3,)).to(data.device)
+    r_pos[..., 1:, [0, 2]] = data[..., :-1, 1:3]
+    '''Add Y-axis rotation to root position'''
+    r_pos = qrot(qinv(r_rot_quat), r_pos)
+
+    r_pos = torch.cumsum(r_pos, dim=-2)
+
+    r_pos[..., 1] = data[..., 3]
+    return r_rot_quat, r_pos
+
+def recover_from_ric(data, joints_num):
+    r_rot_quat, r_pos = recover_root_rot_pos(data)
+    positions = data[..., 4:(joints_num - 1) * 3 + 4]
+    positions = positions.view(positions.shape[:-1] + (-1, 3))
+
+    '''Add Y-axis rotation to local joints'''
+    positions = qrot(qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4,)), positions)
+
+    '''Add root XZ to joints'''
+    positions[..., 0] += r_pos[..., 0:1]
+    positions[..., 2] += r_pos[..., 2:3]
+
+    '''Concate root and joints'''
+    positions = torch.cat([r_pos.unsqueeze(-2), positions], dim=-2)
+
+    return positions
+
+def plot_3d_motion(save_path, kinematic_tree, joints, title, figsize=(10, 10), fps=120, radius=4):
+#     matplotlib.use('Agg')
+
+    title_sp = title.split(' ')
+    if len(title_sp) > 10:
+        title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:])])
+    def init():
+        # ax.set_xlim3d([-radius / 2, radius / 2])
+        # ax.set_ylim3d([0, radius])
+        # ax.set_zlim3d([0, radius])
+        # # print(title)
+        # fig.suptitle(title, fontsize=20)
+        # ax.grid(b=False)
+
+        nb_joints = joints.shape[1]
+        limits = 1000 if nb_joints == 21 else 2
+        ax.set_xlim(-limits, limits)
+        ax.set_ylim(-limits, limits)
+        ax.set_zlim(0, limits)
+        fig.suptitle(title, fontsize=20)
+        ax.grid(False)
+
+    def plot_xzPlane(minx, maxx, miny, minz, maxz):
+        ## Plot a plane XZ
+        verts = [
+            [minx, miny, minz],
+            [minx, miny, maxz],
+            [maxx, miny, maxz],
+            [maxx, miny, minz]
+        ]
+        xz_plane = Poly3DCollection([verts])
+        xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))
+        ax.add_collection3d(xz_plane)
+
+    #         return ax
+
+    # (seq_len, joints_num, 3)
+    data = joints.copy().reshape(len(joints), -1, 3)
+    fig = plt.figure(figsize=figsize)
+    # ax = p3.Axes3D(fig)
+    ax = fig.add_subplot(111, projection='3d')
+    init()
+    MINS = data.min(axis=0).min(axis=0)
+    MAXS = data.max(axis=0).max(axis=0)
+    colors = ['red', 'blue', 'black', 'red', 'blue',
+              'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue',
+             'darkred', 'darkred','darkred','darkred','darkred']
+    frame_number = data.shape[0]
+    #     print(data.shape)
+
+    height_offset = MINS[1]
+    data[:, :, 1] -= height_offset
+    trajec = data[:, 0, [0, 2]]
+
+    data[..., 0] -= data[:, 0:1, 0]
+    data[..., 2] -= data[:, 0:1, 2]
+
+    #     print(trajec.shape)
+
+    def update(index):
+        #         print(index)
+        # ax.lines = []
+        # ax.collections = []
+        for line in list(ax.lines):
+            line.remove()
+        for collection in list(ax.collections):
+            collection.remove()
+        ax.view_init(elev=120, azim=-90)
+        try:
+            ax.dist = 7.5
+        except Exception:
+            pass
+        #         ax =
+        plot_xzPlane(MINS[0]-trajec[index, 0], MAXS[0]-trajec[index, 0], 0, MINS[2]-trajec[index, 1], MAXS[2]-trajec[index, 1])
+#         ax.scatter(data[index, :22, 0], data[index, :22, 1], data[index, :22, 2], color='black', s=3)
+
+        if index > 1:
+            ax.plot3D(trajec[:index, 0]-trajec[index, 0], np.zeros_like(trajec[:index, 0]), trajec[:index, 1]-trajec[index, 1], linewidth=1.0,
+                      color='blue')
+        #             ax = plot_xzPlane(ax, MINS[0], MAXS[0], 0, MINS[2], MAXS[2])
+
+
+        for i, (chain, color) in enumerate(zip(kinematic_tree, colors)):
+#             print(color)
+            if i < 5:
+                linewidth = 4.0
+            else:
+                linewidth = 2.0
+            ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2], linewidth=linewidth, color=color)
+        #         print(trajec[:index, 0].shape)
+
+        plt.axis('off')
+        ax.set_xticklabels([])
+        ax.set_yticklabels([])
+        ax.set_zticklabels([])
+
+    ani = FuncAnimation(fig, update, frames=frame_number, interval=1000/fps, repeat=False)
+
+    ani.save(save_path, fps=fps)
+    plt.close()
+
+
+def plot_3d_motion_ipynb(kinematic_tree, joints, title, figsize=(10, 10), fps=120, radius=4):
+    """
+    在 Jupyter Notebook 中显示 3D 人体运动动画（不保存文件）
+    """
+    title_sp = title.split(' ')
+    if len(title_sp) > 10:
+        title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:])])
+    def init():
+        # ax.set_xlim3d([-radius / 2, radius / 2])
+        # ax.set_ylim3d([0, radius])
+        # ax.set_zlim3d([0, radius])
+        # # print(title)
+        # fig.suptitle(title, fontsize=20)
+        # ax.grid(b=False)
+
+        nb_joints = joints.shape[1]
+        limits = 1000 if nb_joints == 21 else 2
+        ax.set_xlim(-limits, limits)
+        ax.set_ylim(-limits, limits)
+        ax.set_zlim(0, limits)
+        fig.suptitle(title, fontsize=20)
+        ax.grid(b=False)
+
+    def plot_xzPlane(minx, maxx, miny, minz, maxz):
+        ## Plot a plane XZ
+        verts = [
+            [minx, miny, minz],
+            [minx, miny, maxz],
+            [maxx, miny, maxz],
+            [maxx, miny, minz]
+        ]
+        xz_plane = Poly3DCollection([verts])
+        xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))
+        ax.add_collection3d(xz_plane)
+
+    #         return ax
+
+    # (seq_len, joints_num, 3)
+    data = joints.copy().reshape(len(joints), -1, 3)
+    fig = plt.figure(figsize=figsize)
+    # ax = p3.Axes3D(fig)
+    ax = fig.add_subplot(111, projection='3d')
+    init()
+    MINS = data.min(axis=0).min(axis=0)
+    MAXS = data.max(axis=0).max(axis=0)
+    colors = ['red', 'blue', 'black', 'red', 'blue',  
+              'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue',
+             'darkred', 'darkred','darkred','darkred','darkred']
+    frame_number = data.shape[0]
+    #     print(data.shape)
+
+    height_offset = MINS[1]
+    data[:, :, 1] -= height_offset
+    trajec = data[:, 0, [0, 2]]
+    
+    data[..., 0] -= data[:, 0:1, 0]
+    data[..., 2] -= data[:, 0:1, 2]
+
+    #     print(trajec.shape)
+
+    def update(index):
+        #         print(index)
+        # ax.lines = []
+        # ax.collections = []
+        for line in ax.lines:
+            line.remove()
+        for collection in ax.collections:
+            collection.remove()
+        ax.view_init(elev=120, azim=-90)
+        ax.dist = 7.5
+        #         ax =
+        plot_xzPlane(MINS[0]-trajec[index, 0], MAXS[0]-trajec[index, 0], 0, MINS[2]-trajec[index, 1], MAXS[2]-trajec[index, 1])
+#         ax.scatter(data[index, :22, 0], data[index, :22, 1], data[index, :22, 2], color='black', s=3)
+        
+        if index > 1:
+            ax.plot3D(trajec[:index, 0]-trajec[index, 0], np.zeros_like(trajec[:index, 0]), trajec[:index, 1]-trajec[index, 1], linewidth=1.0,
+                      color='blue')
+        #             ax = plot_xzPlane(ax, MINS[0], MAXS[0], 0, MINS[2], MAXS[2])
+        
+        
+        for i, (chain, color) in enumerate(zip(kinematic_tree, colors)):
+#             print(color)
+            if i < 5:
+                linewidth = 4.0
+            else:
+                linewidth = 2.0
+            ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2], linewidth=linewidth, color=color)
+        #         print(trajec[:index, 0].shape)
+
+        plt.axis('off')
+        ax.set_xticklabels([])
+        ax.set_yticklabels([])
+        ax.set_zticklabels([])
+
+    ani = FuncAnimation(fig, update, frames=frame_number, interval=1000/fps, repeat=False)
+    plt.show()

motionvq/paramUtil.py

@@ -0,0 +1,63 @@
+import numpy as np
+
+# Define a kinematic tree for the skeletal struture
+kit_kinematic_chain = [[0, 11, 12, 13, 14, 15], [0, 16, 17, 18, 19, 20], [0, 1, 2, 3, 4], [3, 5, 6, 7], [3, 8, 9, 10]]
+
+kit_raw_offsets = np.array(
+    [
+        [0, 0, 0],
+        [0, 1, 0],
+        [0, 1, 0],
+        [0, 1, 0],
+        [0, 1, 0],
+        [1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [-1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [0, 0, 1],
+        [0, 0, 1],
+        [-1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [0, 0, 1],
+        [0, 0, 1]
+    ]
+)
+
+t2m_raw_offsets = np.array([[0,0,0],
+                           [1,0,0],
+                           [-1,0,0],
+                           [0,1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,1,0],
+                           [0,0,1],
+                           [0,0,1],
+                           [0,1,0],
+                           [1,0,0],
+                           [-1,0,0],
+                           [0,0,1],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0]])
+
+t2m_kinematic_chain = [[0, 2, 5, 8, 11], [0, 1, 4, 7, 10], [0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21], [9, 13, 16, 18, 20]]
+t2m_left_hand_chain = [[20, 22, 23, 24], [20, 34, 35, 36], [20, 25, 26, 27], [20, 31, 32, 33], [20, 28, 29, 30]]
+t2m_right_hand_chain = [[21, 43, 44, 45], [21, 46, 47, 48], [21, 40, 41, 42], [21, 37, 38, 39], [21, 49, 50, 51]]
+
+
+kit_tgt_skel_id = '03950'
+
+t2m_tgt_skel_id = '000021'
+

motionvq/quantize_cnn.py

@@ -0,0 +1,415 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class QuantizeEMAReset(nn.Module):
+    def __init__(self, nb_code, code_dim, args):
+        super().__init__()
+        self.nb_code = nb_code
+        self.code_dim = code_dim
+        self.mu = args.mu
+        self.reset_codebook()
+        
+    def reset_codebook(self):
+        self.init = False
+        self.code_sum = None
+        self.code_count = None
+        # 修复：不强制指定设备，让模型自动处理数据类型和设备
+        self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim))
+
+    def _tile(self, x):
+        nb_code_x, code_dim = x.shape
+        if nb_code_x < self.nb_code:
+            n_repeats = (self.nb_code + nb_code_x - 1) // nb_code_x
+            std = 0.01 / np.sqrt(code_dim)
+            out = x.repeat(n_repeats, 1)
+            out = out + torch.randn_like(out) * std
+        else :
+            out = x
+        return out
+
+    def init_codebook(self, x):
+        out = self._tile(x)
+        self.codebook = out[:self.nb_code]
+        self.code_sum = self.codebook.clone()
+        self.code_count = torch.ones(self.nb_code, device=self.codebook.device)
+        self.init = True
+        
+    @torch.no_grad()
+    def compute_perplexity(self, code_idx) : 
+        # Calculate new centres
+        code_onehot = torch.zeros(self.nb_code, code_idx.shape[0], device=code_idx.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, code_idx.shape[0]), 1)
+
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+        return perplexity
+    
+    @torch.no_grad()
+    def update_codebook(self, x, code_idx):
+        
+        code_onehot = torch.zeros(self.nb_code, x.shape[0], device=x.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, x.shape[0]), 1)
+
+        code_sum = torch.matmul(code_onehot, x)  # nb_code, w
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+
+        out = self._tile(x)
+        code_rand = out[:self.nb_code]
+
+        # Update centres
+        self.code_sum = self.mu * self.code_sum + (1. - self.mu) * code_sum  # w, nb_code
+        self.code_count = self.mu * self.code_count + (1. - self.mu) * code_count  # nb_code
+
+        usage = (self.code_count.view(self.nb_code, 1) >= 1.0).float()
+        code_update = self.code_sum.view(self.nb_code, self.code_dim) / self.code_count.view(self.nb_code, 1)
+
+        self.codebook = usage * code_update + (1 - usage) * code_rand
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+
+            
+        return perplexity
+
+    def preprocess(self, x):
+        # NCT -> NTC -> [NT, C]
+        x = x.permute(0, 2, 1).contiguous()
+        x = x.view(-1, x.shape[-1])  
+        return x
+
+    def quantize(self, x):
+        # Calculate latent code x_l
+        k_w = self.codebook.t()
+        distance = torch.sum(x ** 2, dim=-1, keepdim=True) - 2 * torch.matmul(x, k_w) + torch.sum(k_w ** 2, dim=0,
+                                                                                            keepdim=True)  # (N * L, b)
+        _, code_idx = torch.min(distance, dim=-1)
+        return code_idx
+
+    def dequantize(self, code_idx):
+        x = F.embedding(code_idx, self.codebook)
+        return x
+
+    
+    def forward(self, x):
+        N, width, T = x.shape
+
+        # Preprocess
+        x = self.preprocess(x)
+
+        # Init codebook if not inited
+        if self.training and not self.init:
+            self.init_codebook(x)
+
+        # quantize and dequantize through bottleneck
+        code_idx = self.quantize(x)
+        x_d = self.dequantize(code_idx)
+
+        # Update embeddings
+        if self.training:
+            perplexity = self.update_codebook(x, code_idx)
+        else : 
+            perplexity = self.compute_perplexity(code_idx)
+        
+        # Loss
+        commit_loss = F.mse_loss(x, x_d.detach())
+
+        # Passthrough
+        x_d = x + (x_d - x).detach()
+
+        # Postprocess
+        x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous()   #(N, DIM, T)
+        
+        return x_d, commit_loss, perplexity
+
+
+
+class Quantizer(nn.Module):
+    def __init__(self, n_e, e_dim, beta):
+        super(Quantizer, self).__init__()
+
+        self.e_dim = e_dim
+        self.n_e = n_e
+        self.beta = beta
+
+        self.embedding = nn.Embedding(self.n_e, self.e_dim)
+        self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
+
+    def forward(self, z):
+        
+        N, width, T = z.shape
+        z = self.preprocess(z)
+        assert z.shape[-1] == self.e_dim
+        z_flattened = z.contiguous().view(-1, self.e_dim)
+
+        # B x V
+        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
+            torch.sum(self.embedding.weight**2, dim=1) - 2 * \
+            torch.matmul(z_flattened, self.embedding.weight.t())
+        # B x 1
+        min_encoding_indices = torch.argmin(d, dim=1)
+        z_q = self.embedding(min_encoding_indices).view(z.shape)
+
+        # compute loss for embedding
+        loss = torch.mean((z_q - z.detach())**2) + self.beta * \
+               torch.mean((z_q.detach() - z)**2)
+
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+        z_q = z_q.view(N, T, -1).permute(0, 2, 1).contiguous()   #(N, DIM, T)
+
+        min_encodings = F.one_hot(min_encoding_indices, self.n_e).type(z.dtype)
+        e_mean = torch.mean(min_encodings, dim=0)
+        perplexity = torch.exp(-torch.sum(e_mean*torch.log(e_mean + 1e-10)))
+        return z_q, loss, perplexity
+
+    def quantize(self, z):
+
+        assert z.shape[-1] == self.e_dim
+
+        # B x V
+        d = torch.sum(z ** 2, dim=1, keepdim=True) + \
+            torch.sum(self.embedding.weight ** 2, dim=1) - 2 * \
+            torch.matmul(z, self.embedding.weight.t())
+        # B x 1
+        min_encoding_indices = torch.argmin(d, dim=1)
+        return min_encoding_indices
+
+    def dequantize(self, indices):
+
+        index_flattened = indices.view(-1)
+        z_q = self.embedding(index_flattened)
+        z_q = z_q.view(indices.shape + (self.e_dim, )).contiguous()
+        return z_q
+
+    def preprocess(self, x):
+        # NCT -> NTC -> [NT, C]
+        x = x.permute(0, 2, 1).contiguous()
+        x = x.view(-1, x.shape[-1])  
+        return x
+
+
+
+class QuantizeReset(nn.Module):
+    def __init__(self, nb_code, code_dim, args):
+        super().__init__()
+        self.nb_code = nb_code
+        self.code_dim = code_dim
+        self.reset_codebook()
+        self.codebook = nn.Parameter(torch.randn(nb_code, code_dim))
+        
+    def reset_codebook(self):
+        self.init = False
+        self.code_count = None
+
+    def _tile(self, x):
+        nb_code_x, code_dim = x.shape
+        if nb_code_x < self.nb_code:
+            n_repeats = (self.nb_code + nb_code_x - 1) // nb_code_x
+            std = 0.01 / np.sqrt(code_dim)
+            out = x.repeat(n_repeats, 1)
+            out = out + torch.randn_like(out) * std
+        else :
+            out = x
+        return out
+
+    def init_codebook(self, x):
+        out = self._tile(x)
+        self.codebook = nn.Parameter(out[:self.nb_code])
+        self.code_count = torch.ones(self.nb_code, device=self.codebook.device)
+        self.init = True
+        
+    @torch.no_grad()
+    def compute_perplexity(self, code_idx) : 
+        # Calculate new centres
+        code_onehot = torch.zeros(self.nb_code, code_idx.shape[0], device=code_idx.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, code_idx.shape[0]), 1)
+
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+        return perplexity
+    
+    def update_codebook(self, x, code_idx):
+        
+        code_onehot = torch.zeros(self.nb_code, x.shape[0], device=x.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, x.shape[0]), 1)
+
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+
+        out = self._tile(x)
+        code_rand = out[:self.nb_code]
+
+        # Update centres
+        self.code_count = code_count  # nb_code
+        usage = (self.code_count.view(self.nb_code, 1) >= 1.0).float()
+
+        self.codebook.data = usage * self.codebook.data + (1 - usage) * code_rand
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+
+            
+        return perplexity
+
+    def preprocess(self, x):
+        # NCT -> NTC -> [NT, C]
+        x = x.permute(0, 2, 1).contiguous()
+        x = x.view(-1, x.shape[-1])  
+        return x
+
+    def quantize(self, x):
+        # Calculate latent code x_l
+        k_w = self.codebook.t()
+        distance = torch.sum(x ** 2, dim=-1, keepdim=True) - 2 * torch.matmul(x, k_w) + torch.sum(k_w ** 2, dim=0,
+                                                                                            keepdim=True)  # (N * L, b)
+        _, code_idx = torch.min(distance, dim=-1)
+        return code_idx
+
+    def dequantize(self, code_idx):
+        x = F.embedding(code_idx, self.codebook)
+        return x
+
+    
+    def forward(self, x):
+        N, width, T = x.shape
+        # Preprocess
+        x = self.preprocess(x)
+        # Init codebook if not inited
+        if self.training and not self.init:
+            self.init_codebook(x)
+        # quantize and dequantize through bottleneck
+        code_idx = self.quantize(x)
+        x_d = self.dequantize(code_idx)
+        # Update embeddings
+        if self.training:
+            perplexity = self.update_codebook(x, code_idx)
+        else : 
+            perplexity = self.compute_perplexity(code_idx)
+        
+        # Loss
+        commit_loss = F.mse_loss(x, x_d.detach())
+
+        # Passthrough
+        x_d = x + (x_d - x).detach()
+
+        # Postprocess
+        x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous()   #(N, DIM, T)
+        
+        return x_d, commit_loss, perplexity
+
+    
+class QuantizeEMA(nn.Module):
+    def __init__(self, nb_code, code_dim, args):
+        super().__init__()
+        self.nb_code = nb_code
+        self.code_dim = code_dim
+        self.mu = 0.99
+        self.reset_codebook()
+        
+    def reset_codebook(self):
+        self.init = False
+        self.code_sum = None
+        self.code_count = None
+        # 修复：不强制指定设备，让模型自动处理数据类型和设备
+        self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim))
+
+    def _tile(self, x):
+        nb_code_x, code_dim = x.shape
+        if nb_code_x < self.nb_code:
+            n_repeats = (self.nb_code + nb_code_x - 1) // nb_code_x
+            std = 0.01 / np.sqrt(code_dim)
+            out = x.repeat(n_repeats, 1)
+            out = out + torch.randn_like(out) * std
+        else :
+            out = x
+        return out
+
+    def init_codebook(self, x):
+        out = self._tile(x)
+        self.codebook = out[:self.nb_code]
+        self.code_sum = self.codebook.clone()
+        self.code_count = torch.ones(self.nb_code, device=self.codebook.device)
+        self.init = True
+        
+    @torch.no_grad()
+    def compute_perplexity(self, code_idx) : 
+        # Calculate new centres
+        code_onehot = torch.zeros(self.nb_code, code_idx.shape[0], device=code_idx.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, code_idx.shape[0]), 1)
+
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+        return perplexity
+    
+    @torch.no_grad()
+    def update_codebook(self, x, code_idx):
+        
+        code_onehot = torch.zeros(self.nb_code, x.shape[0], device=x.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, x.shape[0]), 1)
+
+        code_sum = torch.matmul(code_onehot, x)  # nb_code, w
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+
+        # Update centres
+        self.code_sum = self.mu * self.code_sum + (1. - self.mu) * code_sum  # w, nb_code
+        self.code_count = self.mu * self.code_count + (1. - self.mu) * code_count  # nb_code
+
+        code_update = self.code_sum.view(self.nb_code, self.code_dim) / self.code_count.view(self.nb_code, 1)
+
+        self.codebook = code_update
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+            
+        return perplexity
+
+    def preprocess(self, x):
+        # NCT -> NTC -> [NT, C]
+        x = x.permute(0, 2, 1).contiguous()
+        x = x.view(-1, x.shape[-1])  
+        return x
+
+    def quantize(self, x):
+        # Calculate latent code x_l
+        k_w = self.codebook.t()
+        distance = torch.sum(x ** 2, dim=-1, keepdim=True) - 2 * torch.matmul(x, k_w) + torch.sum(k_w ** 2, dim=0,
+                                                                                            keepdim=True)  # (N * L, b)
+        _, code_idx = torch.min(distance, dim=-1)
+        return code_idx
+
+    def dequantize(self, code_idx):
+        x = F.embedding(code_idx, self.codebook)
+        return x
+
+    
+    def forward(self, x):
+        N, width, T = x.shape
+
+        # Preprocess
+        x = self.preprocess(x)
+
+        # Init codebook if not inited
+        if self.training and not self.init:
+            self.init_codebook(x)
+
+        # quantize and dequantize through bottleneck
+        code_idx = self.quantize(x)
+        x_d = self.dequantize(code_idx)
+
+        # Update embeddings
+        if self.training:
+            perplexity = self.update_codebook(x, code_idx)
+        else : 
+            perplexity = self.compute_perplexity(code_idx)
+        
+        # Loss
+        commit_loss = F.mse_loss(x, x_d.detach())
+
+        # Passthrough
+        x_d = x + (x_d - x).detach()
+
+        # Postprocess
+        x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous()   #(N, DIM, T)
+        
+        return x_d, commit_loss, perplexity

motionvq/resnet.py

@@ -0,0 +1,82 @@
+import torch.nn as nn
+import torch
+
+class nonlinearity(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        # swish
+        return x * torch.sigmoid(x)
+
+class ResConv1DBlock(nn.Module):
+    def __init__(self, n_in, n_state, dilation=1, activation='silu', norm=None, dropout=None):
+        super().__init__()
+        padding = dilation
+        self.norm = norm
+        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()
+
+        if activation == "relu":
+            self.activation1 = nn.ReLU()
+            self.activation2 = nn.ReLU()
+            
+        elif activation == "silu":
+            self.activation1 = nonlinearity()
+            self.activation2 = nonlinearity()
+            
+        elif activation == "gelu":
+            self.activation1 = nn.GELU()
+            self.activation2 = nn.GELU()
+            
+        
+
+        self.conv1 = nn.Conv1d(n_in, n_state, 3, 1, padding, dilation)
+        self.conv2 = nn.Conv1d(n_state, n_in, 1, 1, 0,)     
+
+
+    def forward(self, x):
+        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):
+    def __init__(self, n_in, n_depth, dilation_growth_rate=1, reverse_dilation=True, activation='relu', norm=None):
+        super().__init__()
+        
+        blocks = [ResConv1DBlock(n_in, n_in, dilation=dilation_growth_rate ** depth, activation=activation, norm=norm) for depth in range(n_depth)]
+        if reverse_dilation:
+            blocks = blocks[::-1]
+        
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):        
+        return self.model(x)

motionvq/vqvae.py

@@ -0,0 +1,136 @@
+import torch.nn as nn
+from .encdec import Encoder, Decoder
+from .quantize_cnn import QuantizeEMAReset, Quantizer, QuantizeEMA, QuantizeReset
+
+
+class VQVAE_251(nn.Module):
+    def __init__(self,
+                 args,
+                 nb_code=1024,
+                 code_dim=512,
+                 output_emb_width=512,
+                 down_t=3,
+                 stride_t=2,
+                 width=512,
+                 depth=3,
+                 dilation_growth_rate=3,
+                 activation='relu',
+                 norm=None):
+        
+        super().__init__()
+        self.code_dim = code_dim
+        self.num_code = nb_code
+        self.quant = args.quantizer
+        self.encoder = Encoder(251 if args.dataname == 'kit' else 263, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm)
+        self.decoder = Decoder(251 if args.dataname == 'kit' else 263, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm)
+        if args.quantizer == "ema_reset":
+            self.quantizer = QuantizeEMAReset(nb_code, code_dim, args)
+        elif args.quantizer == "orig":
+            self.quantizer = Quantizer(nb_code, code_dim, 1.0)
+        elif args.quantizer == "ema":
+            self.quantizer = QuantizeEMA(nb_code, code_dim, args)
+        elif args.quantizer == "reset":
+            self.quantizer = QuantizeReset(nb_code, code_dim, args)
+
+
+    def preprocess(self, x):
+        # (bs, T, Jx3) -> (bs, Jx3, T)
+        x = x.permute(0,2,1).float()
+        return x
+
+
+    def postprocess(self, x):
+        # (bs, Jx3, T) ->  (bs, T, Jx3)
+        x = x.permute(0,2,1)
+        return x
+
+
+    def encode(self, x):
+        N, T, _ = x.shape
+        x_in = self.preprocess(x)
+        x_encoder = self.encoder(x_in)
+        x_encoder = self.postprocess(x_encoder)
+        x_encoder = x_encoder.contiguous().view(-1, x_encoder.shape[-1])  # (NT, C)
+        code_idx = self.quantizer.quantize(x_encoder)
+        code_idx = code_idx.view(N, -1)
+        return code_idx
+    
+    def get_embeddings(self, x):
+        x_in = self.preprocess(x)
+        x_encoder = self.encoder(x_in)
+        x_quantized, loss, perplexity  = self.quantizer(x_encoder)
+        return x_quantized
+
+
+    def forward(self, x):
+        x_in = self.preprocess(x)
+        # Encode
+        x_encoder = self.encoder(x_in)
+        ## quantization
+        x_quantized, loss, perplexity  = self.quantizer(x_encoder)
+        ## decoder
+        x_decoder = self.decoder(x_quantized)
+        x_out = self.postprocess(x_decoder)
+        return x_out, loss, perplexity
+
+
+    def forward_decoder(self, x):
+        x_d = self.quantizer.dequantize(x)
+        x_d = x_d.view(1, -1, self.code_dim).permute(0, 2, 1).contiguous()
+        
+        # decoder
+        x_decoder = self.decoder(x_d)
+        x_out = self.postprocess(x_decoder)
+        return x_out
+
+    def embeddings_decode(self, x):
+        x_decoder = self.decoder(x)
+        x_out = self.postprocess(x_decoder)
+        return x_out
+
+
+class HumanVQVAE(nn.Module):
+    def __init__(self,
+                 args,
+                 nb_code=512,
+                 code_dim=512,
+                 output_emb_width=512,
+                 down_t=3,
+                 stride_t=2,
+                 width=512,
+                 depth=3,
+                 dilation_growth_rate=3,
+                 activation='relu',
+                 norm=None):
+        
+        super().__init__()
+        
+        self.nb_joints = 21 if args.dataname == 'kit' else 22
+        self.vqvae = VQVAE_251(args, nb_code, code_dim, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm)
+
+    def forward_encoder(self, x):
+        x_out = self.vqvae.forward_encoder(x)
+        return x_out
+
+    def encode(self, x):
+        b, t, c = x.size()
+        quants = self.vqvae.encode(x) # (N, T)
+        return quants
+
+    def forward(self, x):
+
+        x_out, loss, perplexity = self.vqvae(x)
+        
+        return x_out, loss, perplexity
+
+    def forward_decoder(self, x):
+        x_out = self.vqvae.forward_decoder(x)
+        return x_out
+    
+    def get_embeddings(self, x):
+        x_out = self.vqvae.get_embeddings(x)
+        return x_out
+    
+    def embeddings_decode(self, x):
+        x_out = self.vqvae.embeddings_decode(x)
+        return x_out

requirements.txt

@@ -0,0 +1,7 @@
+torch
+transformers==4.56.1
+accelerate
+numpy<2
+matplotlib==3.7.5
+imageio
+imageio-ffmpeg

std.npy

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