initial commit

.gitattributes

@@ -1,35 +1 @@
-*.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

README.md

@@ -1,3 +1,371 @@
 ---
 license: mit
+language:
+- en
+tags:
+- motion
+- human-motion
+- text-to-motion
+- motion-to-text
+- motion-captioning
+- motion-generation
+- autoencoder
+- gpt2
+datasets:
+- humanml3d
+pipeline_tag: text-generation
+library_name: transformers
+base_model:
+- openai-community/gpt2
 ---
+
+# MUGEN
+
+**A Unified Framework for Efficient Motion Understanding and Generation**
+
+MUGEN is a unified motion-language model. One model turns a description into
+human motion, and describes an observed motion in words. Both directions share a
+single motion representation, and that representation is continuous: **no
+codebook, one draw.**
+
+- 📦 Code: [github.com/JYe16/MUGEN](https://github.com/JYe16/MUGEN)
+- 🏋️ This checkpoint: HumanML3D, **K = 2** latent slots, 183.7M parameters
+
+---
+
+## What makes it different
+
+Unified motion-language systems have traditionally coupled generation and
+understanding through a shared *discrete* motion codebook, but quantization
+limits generation quality. The strongest generators buy that quality back at
+growing cost: stacked residual codebooks enlarge the representation, and masked
+decoding stages, long autoregressive rollouts and denoising chains of tens to
+hundreds of steps stretch inference. None of that decoding machinery serves
+understanding.
+
+MUGEN pays neither cost. Generating a motion costs **K language-model steps, one
+draw, and one decoder pass** — with K = 2 here, that is two rollout steps.
+
+## Architecture
+
+```
+MugenForConditionalGeneration                                    183,690,603 params
+├── motion_autoencoder (AdaptiveLengthAutoEncoder)                    48.9M
+│   ├── encoder_backbone   dilated Conv1d ResNet, per-frame features
+│   ├── latent_queries     K learnable queries -> the K slots
+│   ├── encoder_blocks     4 x cross-attention (frames -> slots)
+│   ├── decoder_blocks     4 x cross-attention (slots -> frames)
+│   └── decoder_refiner    dilated Conv1d ResNet -> 263-d motion
+├── language_model (GPT2LMHeadModel)                                  124.4M
+│   └── 12 layers, 768 hidden, vocab 50261 (GPT-2 + <MOT> + 3 reserved)
+├── layer router                                                        9.1M
+│   ├── text_mem_proj      prompt tokens -> routing memory
+│   ├── cross_attn_blocks  2 x cross-attention, per-slot routing queries
+│   ├── static_router_logits   (K, 12) text-independent component
+│   └── layer_router       MLP -> (K, 12) text-conditional component
+├── calibrated latent head                                              0.9M
+│   ├── projector          hidden -> (mu, logvar) per slot
+│   ├── latent_factors     (K*512, 64) global covariance basis U
+│   └── factor_scale_head  hidden -> per-factor log amplitudes
+└── motion_in_projector    latent slots -> GPT-2 embeddings (captioning)  0.4M
+```
+
+**Adaptive-length autoencoder.** Cross-attention compresses a clip of *any*
+length into K = 2 continuous vectors of width 512, and a second cross-attention
+stack expands them back to *any* requested frame count. Decoder queries encode a
+frame's relative phase within the clip rather than an absolute frame index, which
+is why one decoder serves every length, and why you can decode the same latents
+at 60 frames and at 200 frames.
+
+**Depth-routed hidden states.** A text-conditioned router gives each latent slot
+its own soft mixture over all twelve transformer layers, so a slot reads from the
+depth it needs instead of squeezing every piece of motion evidence through the
+final layer. Both halves of the routing logit are tanh-bounded, so no logit
+margin can saturate the routing softmax.
+
+**Calibrated latent head.** The head predicts
+`N(mu, U diag(a)^2 U^T + diag(sigma^2))` over the whole flattened `K x 512`
+latent, with rank 64, trained by exact maximum likelihood. One draw therefore
+carries text-conditional variance that is *correlated across slots*, which is
+what a single-step sampler has to supply all at once.
+
+| | |
+|---|---|
+| Motion features | HumanML3D, 263-d per frame, 20 fps, 22 joints |
+| Latent slots | K = 2, width 512 |
+| Language model | GPT-2 (124M), fine-tuned |
+| Covariance rank | 64 |
+| Training | Joint generation + understanding, 240 epochs (Stage 2) |
+| Parameters | 183,690,603 |
+| Precision | float32 |
+
+---
+
+## Installation
+
+```bash
+pip install torch transformers safetensors numpy
+```
+
+`trust_remote_code=True` is required: the autoencoder, the router and the
+calibrated head are custom modules that ship with this repository.
+
+**Nothing else is needed.** The model carries its own tokenizer and its own
+HumanML3D feature statistics, so it returns motion in real units without you
+downloading the dataset.
+
+---
+
+## Quick start
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "zy22b/MUGEN", trust_remote_code=True
+).eval()
+tokenizer = AutoTokenizer.from_pretrained("zy22b/MUGEN")
+
+features = model.generate_motion(
+    ["a person walks forward and then waves with the right hand."],
+    lengths=[120],              # frames at 20 fps, so 6 seconds
+    tokenizer=tokenizer,
+)
+print(features.shape)                       # torch.Size([1, 120, 263])
+
+joints = model.features_to_joints(features)
+print(joints.shape)                         # torch.Size([1, 120, 22, 3]) metres
+
+print(model.generate_caption(features, tokenizer=tokenizer))
+# ['a person walks forward, turns to the left, and walks back.']
+```
+
+That last line captions the motion the model just generated, so it changes from
+call to call: the motion is *sampled*, not a fixed function of the prompt.
+
+---
+
+## Inference tutorial
+
+### 1. Text to motion
+
+```python
+prompts = [
+    "a person walks forward and then waves with the right hand.",
+    "the person jumps up and lands with both feet together.",
+    "a man slowly sits down on a chair.",
+    "a person runs in a circle to the left.",
+]
+
+features = model.generate_motion(prompts, lengths=120, tokenizer=tokenizer)
+```
+
+`generate_motion` returns `(batch, frames, 263)` denormalised HumanML3D features.
+
+| Argument | Meaning |
+|---|---|
+| `texts` | one description, or a list of them |
+| `lengths` | frames per description at 20 fps; an `int` applies to the whole batch |
+| `temperature` | multiplier on the sampled perturbation; defaults to the calibrated `1.0` |
+| `generator` | a `torch.Generator` for a reproducible draw |
+| `denormalize` | `True` (default) returns raw HumanML3D units |
+
+**Choosing a length.** HumanML3D clips are multiples of 4 frames and at most 196
+(9.8 s), so stay in `[20, 196]` and on the multiple-of-4 grid to stay in
+distribution. The decoder will happily produce other lengths, but you are then
+outside what it was trained on.
+
+**Per-row lengths.** Because decoder queries are relative-phase, a row must be
+decoded at its own length rather than cropped from a longer one. Pass a list and
+the model handles it; the returned tensor is zero-padded to the longest row.
+
+```python
+features = model.generate_motion(prompts[:2], lengths=[60, 196], tokenizer=tokenizer)
+# (2, 196, 263); row 0 is valid up to frame 60
+```
+
+### 2. Temperature, and why the default is 1.0
+
+The head predicts a distribution, and `temperature` scales the whole zero-mean
+perturbation around its mean.
+
+```python
+# The calibrated conditional distribution. This is the protocol every reported
+# number uses, and different calls give different motions for the same prompt.
+features = model.generate_motion(prompts, lengths=120, tokenizer=tokenizer)
+
+# The distribution mean. Deterministic, and a *regression* protocol: numbers
+# obtained this way are not comparable to generative ones.
+mean_motion = model.generate_motion(prompts, lengths=120, tokenizer=tokenizer,
+                                    temperature=0.0)
+
+# Reproducible sampling.
+g = torch.Generator().manual_seed(1234)
+features = model.generate_motion(prompts, lengths=120, tokenizer=tokenizer, generator=g)
+```
+
+Sampling at temperature 1.0 is what makes the same prompt yield genuinely
+different, plausible motions. Captioning five seeded draws of the same prompt
+shows how far apart they land:
+
+```
+prompt: "a person walks forward and then waves with the right hand."
+  seed 0  a person walks forward and then turns around and walks back.
+  seed 1  a person walks forward and reaches out with their right hand.
+  seed 2  a person walks forward and then waves with their left hand.
+  seed 3  a person walks forward and puts their right hand on their head.
+  seed 4  a person walks forward and then puts something on a counter.
+```
+
+Lowering the temperature trades that diversity for proximity to the conditional
+mean.
+
+### 3. Motion to text
+
+```python
+import numpy as np
+
+motion = np.load("000004.npy")                      # (T, 263) HumanML3D features
+clip = torch.from_numpy(motion).float().unsqueeze(0)
+
+print(model.generate_caption(clip, tokenizer=tokenizer))
+# ['a person bends over to touch the ground with their hands, then stands up straight.']
+```
+
+Pass `normalized=True` if your features are already standardised by the dataset
+statistics. Captioning is greedy by default; `max_new_tokens` and `num_beams`
+override the config values.
+
+### 4. The shared latent interface
+
+The K slots are the only motion representation in the system: the language model
+generates them for text-to-motion and reads these same vectors back for
+captioning. You can work with them directly.
+
+```python
+latents = model.encode_motion(clip)          # (1, 2, 512): a whole clip, two vectors
+print(latents.shape)
+
+# Decode at the original length, and at a different one.
+same = model.decode_motion(latents, clip.shape[1])
+longer = model.decode_motion(latents, 200)
+
+# Caption straight from latents, skipping the encoder.
+print(model.generate_caption(latents=latents, tokenizer=tokenizer))
+
+# Text -> latents, without decoding.
+z = model.text_to_latents(["a person waves."], tokenizer=tokenizer)
+```
+
+Interpolating between two clips is then a two-line operation:
+
+```python
+a = model.encode_motion(clip_a)
+b = model.encode_motion(clip_b)
+blend = model.decode_motion(0.5 * a + 0.5 * b, 120)
+```
+
+### 5. Running on GPU, and in batches
+
+```python
+model = model.to("cuda")
+features = model.generate_motion(prompts, lengths=120, tokenizer=tokenizer)  # follows the model
+```
+
+Batched prompts are left-padded, which reproduces the training and evaluation
+behaviour exactly. Single-prompt calls involve no padding at all.
+
+### 6. Saving and rendering
+
+`features_to_joints` gives `(batch, frames, 22, 3)` positions in metres, on the
+HumanML3D skeleton, with the y axis vertical.
+
+```python
+joints = model.features_to_joints(features)
+np.save("sample.npy", joints[0].cpu().numpy())
+```
+
+For animation, feed those positions to any HumanML3D-compatible renderer. The
+[code repository](https://github.com/JYe16/MUGEN) includes
+`motGPT/utils/render_utils.py`, which writes an mp4 directly from a `[T, J, 3]`
+array.
+
+### 7. Sanity check
+
+If you swap in different feature statistics, the failure is silent: shapes stay
+right and the numbers look plausible while the whole skeleton is mis-scaled. One
+assertion catches it.
+
+```python
+joints = model.features_to_joints(features)
+high, low = float(joints[..., 1].max()), float(joints[..., 1].min())
+assert 1.0 < high < 2.5 and -0.5 < low < 0.5, "implausible body height"
+```
+
+---
+
+## Reproducibility
+
+The inference path in `modeling_mugen.py` was verified against the original
+training code on the same checkpoint. Routed hidden states, `mu`, `logvar`, the
+factor log-amplitudes, the autoencoder encode and decode outputs, and a
+seeded draw all agree to **`max|diff| = 0.000e+00`**, and generated captions match
+token for token.
+
+Two consequences worth knowing:
+
+- A given seed reproduces the same draw here as it does in the research code,
+  because the factor noise is drawn before the diagonal noise in both.
+- `router_eval_tau` (1.5) is the converged end of the training temperature
+  anneal. Changing it changes which transformer depth each slot reads from, so
+  leave it alone unless you are deliberately studying the router.
+
+Tested with `transformers` 4.47 and PyTorch 2.x. The rollout passes explicit
+`position_ids`, so behaviour does not drift with the transformers version.
+
+## Limitations
+
+- Trained on **HumanML3D** only: everyday single-person motion, English
+  descriptions, 20 fps, 22 joints. Multi-person interaction, object manipulation,
+  and highly stylised motion are out of distribution.
+- Length is an **input**, not something the model infers from the text. Ask for
+  120 frames and you get 120 frames whether or not the description warrants it.
+- K = 2 slots is a deliberately tight budget, chosen for this dataset. It is not
+  a universal setting: the SnapMoGen model in the code repository uses K = 4, and
+  the best budget is dataset-dependent.
+- Captions come from a fine-tuned GPT-2 (124M) and are short and generic
+  relative to human references. They describe the dominant action reliably and
+  fine detail unreliably. On HumanML3D clip `000004`, whose reference reads *"a
+  man is pretending to be a chicken, constantly pecking at the ground and waving
+  his arms like a chicken"*, the model produces *"a person bends over to touch
+  the ground with their hands, then stands up straight"*: the gross movement is
+  right, the intent and the arm motion are gone.
+- Not a safety-filtered model. It generates skeletal motion and short English
+  captions, and it has no content filtering of any kind.
+
+## Citation
+
+The paper is under review. A citation entry will be added here once it appears;
+in the meantime please cite the repository and this model card.
+
+```bibtex
+@misc{mugen2026,
+  title  = {MUGEN: A Unified Framework for Efficient Motion Understanding and Generation},
+  note   = {Under review. Code: https://github.com/JYe16/MUGEN,
+            model: https://huggingface.co/zy22b/MUGEN},
+  year   = {2026}
+}
+```
+
+## Acknowledgements
+
+Built on [HumanML3D](https://github.com/EricGuo5513/HumanML3D),
+[MotionGPT](https://github.com/OpenMotionLab/MotionGPT),
+[MLD](https://github.com/ChenFengYe/motion-latent-diffusion), and our earlier
+discrete-codebook system
+[GeoMotionGPT](https://huggingface.co/zy22b/GeoMotionGPT).
+
+## License
+
+MIT

added_tokens.json

@@ -0,0 +1,6 @@
+{
+  "<MOT>": 50260,
+  "<motion_id_0>": 50257,
+  "<motion_id_1>": 50258,
+  "<motion_id_2>": 50259
+}

config.json

@@ -0,0 +1,58 @@
+{
+  "activation_function": "gelu_new",
+  "alae_activation": "gelu",
+  "alae_depth": 3,
+  "alae_dilation_growth_rate": 3,
+  "alae_dim_feedforward": 2048,
+  "alae_dropout": 0.05,
+  "alae_hidden_dim": 512,
+  "alae_max_decode_len": 256,
+  "alae_nhead": 8,
+  "alae_norm": null,
+  "alae_num_decoder_layers": 4,
+  "alae_num_encoder_layers": 4,
+  "alae_num_res_blocks": 2,
+  "architectures": [
+    "MugenForConditionalGeneration"
+  ],
+  "attn_pdrop": 0.1,
+  "auto_map": {
+    "AutoConfig": "configuration_mugen.MugenConfig",
+    "AutoModelForCausalLM": "modeling_mugen.MugenForConditionalGeneration"
+  },
+  "bos_token_id": 50256,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 50256,
+  "eval_sample_temperature": 1.0,
+  "fps": 20,
+  "initializer_range": 0.02,
+  "k_latent_slots": 2,
+  "latent_dim": 512,
+  "latent_low_rank": 64,
+  "layer_norm_epsilon": 1e-05,
+  "length_multiple": 4,
+  "m2t_max_new_tokens": 64,
+  "m2t_num_beams": 1,
+  "m2t_prompt_prefix": "Please describe the following human motion using plain text:",
+  "model_type": "mugen",
+  "mot_token": "<MOT>",
+  "mot_token_id": 50260,
+  "motion_input_dim": 263,
+  "n_embd": 768,
+  "n_head": 12,
+  "n_inner": null,
+  "n_layer": 12,
+  "n_positions": 1024,
+  "num_cross_attn_layers": 2,
+  "num_joints": 22,
+  "pad_token_id": 50256,
+  "resid_pdrop": 0.1,
+  "router_delta_scale": 4.0,
+  "router_eval_tau": 1.5,
+  "router_hidden": 512,
+  "router_static_scale": 4.0,
+  "t2m_prompt_template": "Please generate human motion based on the following textual description: {text} {mot}",
+  "torch_dtype": "float32",
+  "transformers_version": "4.51.0",
+  "vocab_size": 50261
+}

configuration_mugen.py

@@ -0,0 +1,231 @@
+"""MUGEN configuration.
+
+MUGEN couples an adaptive-length motion autoencoder (ALAE) with a GPT-2
+language model. The autoencoder compresses a motion clip of any length into K
+continuous latent slots, and those slots are the only motion representation in
+the system: the language model generates them for text-to-motion and reads
+them back for motion understanding (captioning).
+"""
+
+from transformers import PretrainedConfig
+
+
+class MugenConfig(PretrainedConfig):
+    """Configuration for :class:`MugenForConditionalGeneration`.
+
+    Args:
+        motion_input_dim (`int`, *optional*, defaults to 263):
+            Dimensionality of one motion frame. 263 is the HumanML3D feature
+            layout (root velocities, RIC joint positions, 6D rotations, foot
+            contacts).
+        k_latent_slots (`int`, *optional*, defaults to 2):
+            Number of continuous latent slots K. The whole clip, whatever its
+            length, is represented by exactly K vectors.
+        latent_dim (`int`, *optional*, defaults to 512):
+            Width of one latent slot.
+        alae_hidden_dim (`int`, *optional*, defaults to 512):
+            Width of the autoencoder's convolutional trunk.
+        alae_depth (`int`, *optional*, defaults to 3):
+            Number of dilated convolutions inside one residual block.
+        alae_dilation_growth_rate (`int`, *optional*, defaults to 3):
+            Dilation growth factor across the residual stack.
+        alae_activation (`str`, *optional*, defaults to `"gelu"`):
+            Activation used throughout the autoencoder.
+        alae_norm (`str`, *optional*, defaults to `None`):
+            Normalisation inside the residual blocks (`None`, `"LN"`, `"GN"`,
+            `"BN"`).
+        alae_num_res_blocks (`int`, *optional*, defaults to 2):
+            Residual blocks in the encoder trunk and the decoder refiner.
+        alae_num_encoder_layers (`int`, *optional*, defaults to 4):
+            Cross-attention blocks that compress frames into the K slots.
+        alae_num_decoder_layers (`int`, *optional*, defaults to 4):
+            Cross-attention blocks that expand the K slots back to frames.
+        alae_nhead (`int`, *optional*, defaults to 8):
+            Attention heads in the autoencoder's cross-attention blocks.
+        alae_dim_feedforward (`int`, *optional*, defaults to 2048):
+            Feed-forward width in those blocks.
+        alae_dropout (`float`, *optional*, defaults to 0.05):
+            Dropout in those blocks (inactive at inference).
+        alae_max_decode_len (`int`, *optional*, defaults to 256):
+            Size of the decoder's positional table. Query positions are spread
+            across the whole table so a position encodes the frame's relative
+            phase in the clip rather than an absolute frame index, which is
+            what lets one decoder serve every clip length.
+        vocab_size (`int`, *optional*, defaults to 50261):
+            Language-model vocabulary: GPT-2's 50257 plus three reserved
+            motion-id slots and the `<MOT>` seed token.
+        n_positions, n_embd, n_layer, n_head, n_inner, activation_function,
+        resid_pdrop, embd_pdrop, attn_pdrop, layer_norm_epsilon,
+        initializer_range:
+            Standard GPT-2 hyperparameters.
+        num_cross_attn_layers (`int`, *optional*, defaults to 2):
+            Cross-attention blocks in the layer router's text encoder.
+        router_hidden (`int`, *optional*, defaults to 512):
+            Hidden width of the router MLP.
+        router_static_scale (`float`, *optional*, defaults to 4.0):
+            Bound on the router's text-independent logit component.
+        router_delta_scale (`float`, *optional*, defaults to 4.0):
+            Bound on the router's text-conditional logit component. Both parts
+            are tanh-capped, so no logit margin can saturate the softmax.
+        router_eval_tau (`float`, *optional*, defaults to 1.5):
+            Inference temperature of the routing softmax. This is the converged
+            end of the training anneal; changing it changes which transformer
+            depth each slot reads from.
+        latent_low_rank (`int`, *optional*, defaults to 64):
+            Rank r of the calibrated head's covariance. The head predicts
+            `N(mu, U diag(a)^2 U^T + diag(sigma^2))` over the flattened K*D
+            latent, so a single draw carries variance correlated across slots.
+        eval_sample_temperature (`float`, *optional*, defaults to 1.0):
+            Default multiplier on the sampled perturbation around `mu`. 1.0 is
+            the calibrated conditional distribution and the setting every
+            reported number uses; 0.0 decodes the mean.
+        m2t_max_new_tokens (`int`, *optional*, defaults to 64):
+            Caption length cap for motion-to-text generation.
+        m2t_num_beams (`int`, *optional*, defaults to 1):
+            Beam count for captioning.
+        mot_token (`str`, *optional*, defaults to `"<MOT>"`):
+            Seed token that ends the text-to-motion prompt.
+        mot_token_id (`int`, *optional*, defaults to 50260):
+            Token id of `mot_token` in the shipped tokenizer.
+        t2m_prompt_template (`str`, *optional*):
+            Prompt used for text-to-motion. Must end in `mot_token`.
+        m2t_prompt_prefix (`str`, *optional*):
+            Instruction placed before the motion slots for captioning.
+        fps (`int`, *optional*, defaults to 20):
+            Frame rate of the HumanML3D features this model was trained on.
+        num_joints (`int`, *optional*, defaults to 22):
+            Joint count recovered by `features_to_joints`.
+        length_multiple (`int`, *optional*, defaults to 4):
+            HumanML3D crops clips to a multiple of 4 frames; requesting a
+            length off this grid moves you away from the training distribution.
+
+    Example:
+        ```python
+        from transformers import AutoConfig
+
+        config = AutoConfig.from_pretrained("zy22b/MUGEN", trust_remote_code=True)
+        print(config.k_latent_slots)  # 2
+        ```
+    """
+
+    model_type = "mugen"
+
+    def __init__(
+        self,
+        # ---- motion autoencoder (ALAE) ----
+        motion_input_dim: int = 263,
+        k_latent_slots: int = 2,
+        latent_dim: int = 512,
+        alae_hidden_dim: int = 512,
+        alae_depth: int = 3,
+        alae_dilation_growth_rate: int = 3,
+        alae_activation: str = "gelu",
+        alae_norm=None,
+        alae_num_res_blocks: int = 2,
+        alae_num_encoder_layers: int = 4,
+        alae_num_decoder_layers: int = 4,
+        alae_nhead: int = 8,
+        alae_dim_feedforward: int = 2048,
+        alae_dropout: float = 0.05,
+        alae_max_decode_len: int = 256,
+        # ---- language model (GPT-2) ----
+        vocab_size: int = 50261,
+        n_positions: int = 1024,
+        n_embd: int = 768,
+        n_layer: int = 12,
+        n_head: int = 12,
+        n_inner=None,
+        activation_function: str = "gelu_new",
+        resid_pdrop: float = 0.1,
+        embd_pdrop: float = 0.1,
+        attn_pdrop: float = 0.1,
+        layer_norm_epsilon: float = 1e-5,
+        initializer_range: float = 0.02,
+        # ---- layer router ----
+        num_cross_attn_layers: int = 2,
+        router_hidden: int = 512,
+        router_static_scale: float = 4.0,
+        router_delta_scale: float = 4.0,
+        router_eval_tau: float = 1.5,
+        # ---- calibrated latent head ----
+        latent_low_rank: int = 64,
+        eval_sample_temperature: float = 1.0,
+        # ---- motion understanding ----
+        m2t_max_new_tokens: int = 64,
+        m2t_num_beams: int = 1,
+        # ---- prompting ----
+        mot_token: str = "<MOT>",
+        mot_token_id: int = 50260,
+        t2m_prompt_template: str = (
+            "Please generate human motion based on the following textual "
+            "description: {text} {mot}"
+        ),
+        m2t_prompt_prefix: str = (
+            "Please describe the following human motion using plain text:"
+        ),
+        # ---- data conventions ----
+        fps: int = 20,
+        num_joints: int = 22,
+        length_multiple: int = 4,
+        # ---- special tokens ----
+        bos_token_id: int = 50256,
+        eos_token_id: int = 50256,
+        pad_token_id: int = 50256,
+        **kwargs,
+    ):
+        self.motion_input_dim = motion_input_dim
+        self.k_latent_slots = k_latent_slots
+        self.latent_dim = latent_dim
+        self.alae_hidden_dim = alae_hidden_dim
+        self.alae_depth = alae_depth
+        self.alae_dilation_growth_rate = alae_dilation_growth_rate
+        self.alae_activation = alae_activation
+        self.alae_norm = alae_norm
+        self.alae_num_res_blocks = alae_num_res_blocks
+        self.alae_num_encoder_layers = alae_num_encoder_layers
+        self.alae_num_decoder_layers = alae_num_decoder_layers
+        self.alae_nhead = alae_nhead
+        self.alae_dim_feedforward = alae_dim_feedforward
+        self.alae_dropout = alae_dropout
+        self.alae_max_decode_len = alae_max_decode_len
+
+        self.vocab_size = vocab_size
+        self.n_positions = n_positions
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_inner = n_inner
+        self.activation_function = activation_function
+        self.resid_pdrop = resid_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.attn_pdrop = attn_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+
+        self.num_cross_attn_layers = num_cross_attn_layers
+        self.router_hidden = router_hidden
+        self.router_static_scale = router_static_scale
+        self.router_delta_scale = router_delta_scale
+        self.router_eval_tau = router_eval_tau
+
+        self.latent_low_rank = latent_low_rank
+        self.eval_sample_temperature = eval_sample_temperature
+
+        self.m2t_max_new_tokens = m2t_max_new_tokens
+        self.m2t_num_beams = m2t_num_beams
+
+        self.mot_token = mot_token
+        self.mot_token_id = mot_token_id
+        self.t2m_prompt_template = t2m_prompt_template
+        self.m2t_prompt_prefix = m2t_prompt_prefix
+
+        self.fps = fps
+        self.num_joints = num_joints
+        self.length_multiple = length_multiple
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            **kwargs,
+        )

model.safetensors

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

modeling_mugen.py

@@ -0,0 +1,899 @@
+"""MUGEN: a unified motion-language model.
+
+The model has three parts:
+
+1. An adaptive-length motion autoencoder (ALAE). Cross-attention compresses a
+   clip of any length into K continuous latent slots, and a second
+   cross-attention stack expands those slots back to any requested number of
+   frames. There is no codebook and no quantisation anywhere in the pipeline.
+2. A GPT-2 language model. For text-to-motion it rolls out K hidden states
+   after a `<MOT>` seed token; for captioning it reads the K slots back as
+   continuous input embeddings.
+3. A layer router and a calibrated latent head. The router gives each slot its
+   own soft mixture over all twelve transformer layers, so a slot reads from
+   the depth it needs instead of the final layer only. The head predicts a
+   joint Gaussian over the whole flattened latent set with a rank-64 plus
+   diagonal covariance, so one draw carries text-conditional variance that is
+   correlated across slots.
+
+Generating a motion therefore costs K language-model steps, one draw, and one
+decoder pass. No iterative refinement, no residual stages, no denoising chain.
+
+Usage:
+    ```python
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    model = AutoModelForCausalLM.from_pretrained("zy22b/MUGEN", trust_remote_code=True).eval()
+    tokenizer = AutoTokenizer.from_pretrained("zy22b/MUGEN")
+
+    feats = model.generate_motion(["a person walks forward and waves."],
+                                  lengths=[120], tokenizer=tokenizer)
+    joints = model.features_to_joints(feats)          # (1, 120, 22, 3)
+    caption = model.generate_caption(feats, tokenizer=tokenizer)
+    ```
+"""
+
+from __future__ import annotations
+
+import math
+from typing import List, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GPT2Config, GPT2LMHeadModel, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+try:
+    from .configuration_mugen import MugenConfig
+except ImportError:  # standalone / trust_remote_code loading
+    from configuration_mugen import MugenConfig
+
+
+# =====================================================================
+# Motion autoencoder building blocks
+# =====================================================================
+
+class Swish(nn.Module):
+    """x * sigmoid(x)."""
+
+    def forward(self, x):
+        return x * torch.sigmoid(x)
+
+
+def _activation(name: str) -> nn.Module:
+    if name == "relu":
+        return nn.ReLU()
+    if name == "gelu":
+        return nn.GELU()
+    if name == "silu":
+        return Swish()
+    raise ValueError(f"Unsupported activation: {name}")
+
+
+class ResConv1DBlock(nn.Module):
+    """Dilated residual convolution over the time axis."""
+
+    def __init__(self, n_in, n_state, dilation=1, activation="gelu", norm=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(32, n_in, eps=1e-6, affine=True)
+            self.norm2 = nn.GroupNorm(32, n_in, eps=1e-6, affine=True)
+        elif norm == "BN":
+            self.norm1 = nn.BatchNorm1d(n_in, eps=1e-6, affine=True)
+            self.norm2 = nn.BatchNorm1d(n_in, eps=1e-6, affine=True)
+        else:
+            self.norm1 = nn.Identity()
+            self.norm2 = nn.Identity()
+
+        self.activation1 = _activation(activation)
+        self.activation2 = _activation(activation)
+        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)
+        return x + x_orig
+
+
+class Resnet1D(nn.Module):
+    """Stack of dilated residual convolutions."""
+
+    def __init__(self, n_in, n_depth, dilation_growth_rate=1,
+                 reverse_dilation=True, activation="gelu", 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)
+
+
+class CrossAttentionBlock(nn.Module):
+    """Pre-norm self-attention, cross-attention and feed-forward block.
+
+    `memory_key_padding_mask` (True marks padding) is forwarded to the
+    cross-attention. Passing `None` reproduces the unmasked block exactly.
+    """
+
+    def __init__(self, d_model, nhead, dim_feedforward, dropout=0.1, activation="gelu"):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.drop1 = nn.Dropout(dropout)
+        self.drop2 = nn.Dropout(dropout)
+        self.drop3 = nn.Dropout(dropout)
+        self.activation = _activation(activation)
+
+    def forward(self, target, memory, memory_key_padding_mask=None):
+        x = self.norm1(target)
+        x_sa, _ = self.self_attn(x, x, x, need_weights=False)
+        target = target + self.drop1(x_sa)
+
+        x = self.norm2(target)
+        x_ca, _ = self.cross_attn(
+            x, memory, memory, need_weights=False,
+            key_padding_mask=memory_key_padding_mask,
+        )
+        target = target + self.drop2(x_ca)
+
+        x = self.norm3(target)
+        x_ff = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        target = target + self.drop3(x_ff)
+        return target
+
+
+class MotionEncoderBackbone(nn.Module):
+    """Per-frame convolutional trunk feeding the cross-attention encoder."""
+
+    def __init__(self, input_dim, hidden_dim, latent_dim, depth,
+                 dilation_growth_rate, activation, norm, num_res_blocks):
+        super().__init__()
+        blocks = [
+            nn.Conv1d(input_dim, hidden_dim, kernel_size=3, stride=1, padding=1),
+            _activation(activation),
+        ]
+        for _ in range(num_res_blocks):
+            blocks.append(Resnet1D(hidden_dim, depth, dilation_growth_rate,
+                                   reverse_dilation=False, activation=activation, norm=norm))
+        self.backbone = nn.Sequential(*blocks)
+        self.output_proj = nn.Conv1d(hidden_dim, latent_dim, kernel_size=1)
+
+    def forward(self, x):
+        return self.output_proj(self.backbone(x))
+
+
+class MotionDecoderRefiner(nn.Module):
+    """Convolutional refiner turning decoded per-frame features into motion."""
+
+    def __init__(self, latent_dim, hidden_dim, output_dim, depth,
+                 dilation_growth_rate, activation, norm, num_res_blocks):
+        super().__init__()
+        blocks = [
+            nn.Conv1d(latent_dim, hidden_dim, kernel_size=3, stride=1, padding=1),
+            _activation(activation),
+        ]
+        for _ in range(num_res_blocks):
+            blocks.append(Resnet1D(hidden_dim, depth, dilation_growth_rate,
+                                   reverse_dilation=True, activation=activation, norm=norm))
+        blocks.extend([
+            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, stride=1, padding=1),
+            _activation(activation),
+            nn.Conv1d(hidden_dim, output_dim, kernel_size=3, stride=1, padding=1),
+        ])
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+def build_sine_position_encoding(length, dim, device, dtype=torch.float32):
+    pe = torch.zeros(length, dim, device=device, dtype=dtype)
+    position = torch.arange(length, device=device, dtype=dtype).unsqueeze(1)
+    div_term = torch.exp(
+        torch.arange(0, dim, 2, device=device, dtype=dtype) * (-math.log(10000.0) / dim)
+    )
+    pe[:, 0::2] = torch.sin(position * div_term)
+    pe[:, 1::2] = torch.cos(position * div_term)
+    return pe
+
+
+class AdaptiveLengthAutoEncoder(nn.Module):
+    """Compress a clip of any length into K latent slots, and back again.
+
+    `encode` returns `(batch, K, latent_dim)` regardless of the input length.
+    `decode` takes those slots plus a target frame count. Decoder query
+    positions are spread evenly across the whole positional table, so a query
+    encodes a frame's relative phase within the clip rather than an absolute
+    frame index. That is what lets one decoder serve every clip length.
+    """
+
+    def __init__(self, config: MugenConfig):
+        super().__init__()
+        self.input_dim = config.motion_input_dim
+        self.k = config.k_latent_slots
+        self.latent_dim = config.latent_dim
+        self.max_decode_len = config.alae_max_decode_len
+
+        common = dict(
+            depth=config.alae_depth,
+            dilation_growth_rate=config.alae_dilation_growth_rate,
+            activation=config.alae_activation,
+            norm=config.alae_norm,
+            num_res_blocks=config.alae_num_res_blocks,
+        )
+        self.encoder_backbone = MotionEncoderBackbone(
+            input_dim=config.motion_input_dim,
+            hidden_dim=config.alae_hidden_dim,
+            latent_dim=config.latent_dim,
+            **common,
+        )
+        self.latent_queries = nn.Parameter(torch.randn(self.k, config.latent_dim) * 0.02)
+
+        block_kwargs = dict(
+            d_model=config.latent_dim,
+            nhead=config.alae_nhead,
+            dim_feedforward=config.alae_dim_feedforward,
+            dropout=config.alae_dropout,
+            activation=config.alae_activation,
+        )
+        self.encoder_blocks = nn.ModuleList(
+            [CrossAttentionBlock(**block_kwargs) for _ in range(config.alae_num_encoder_layers)]
+        )
+        self.encoder_norm = nn.LayerNorm(config.latent_dim)
+
+        self.decoder_blocks = nn.ModuleList(
+            [CrossAttentionBlock(**block_kwargs) for _ in range(config.alae_num_decoder_layers)]
+        )
+        self.decoder_norm = nn.LayerNorm(config.latent_dim)
+        self.decoder_refiner = MotionDecoderRefiner(
+            latent_dim=config.latent_dim,
+            hidden_dim=config.alae_hidden_dim,
+            output_dim=config.motion_input_dim,
+            **common,
+        )
+
+        self.register_buffer(
+            "decoder_query_pe",
+            build_sine_position_encoding(self.max_decode_len, config.latent_dim,
+                                         device=torch.device("cpu")),
+            persistent=False,
+        )
+
+    def ensure_decode_capacity(self, target_len: int):
+        target_len = int(target_len)
+        if target_len <= self.max_decode_len:
+            return
+        self.max_decode_len = target_len
+        self.register_buffer(
+            "decoder_query_pe",
+            build_sine_position_encoding(target_len, self.latent_dim,
+                                         device=torch.device("cpu")),
+            persistent=False,
+        )
+
+    def encode(self, motion: torch.Tensor) -> torch.Tensor:
+        """`(batch, time, input_dim)` -> `(batch, K, latent_dim)`."""
+        x_in = motion.permute(0, 2, 1).float()
+        memory = self.encoder_backbone(x_in).permute(0, 2, 1)
+        queries = self.latent_queries.unsqueeze(0).expand(memory.shape[0], -1, -1)
+        latents = queries
+        for block in self.encoder_blocks:
+            latents = block(latents, memory)
+        return self.encoder_norm(latents)
+
+    def decode(self, latents: torch.Tensor, target_len: int) -> torch.Tensor:
+        """`(batch, K, latent_dim)` -> `(batch, target_len, input_dim)`."""
+        target_len = int(target_len)
+        self.ensure_decode_capacity(target_len)
+        if target_len > 1:
+            idx = torch.linspace(0, self.max_decode_len - 1, target_len,
+                                 device=latents.device).round().long()
+        else:
+            idx = torch.zeros(target_len, device=latents.device, dtype=torch.long)
+        query_pe = self.decoder_query_pe.to(device=latents.device, dtype=latents.dtype)[idx]
+        decoded = query_pe.unsqueeze(0).expand(latents.shape[0], -1, -1)
+        for block in self.decoder_blocks:
+            decoded = block(decoded, latents)
+        decoded = self.decoder_norm(decoded).permute(0, 2, 1)
+        return self.decoder_refiner(decoded).permute(0, 2, 1)
+
+    def forward(self, motion: torch.Tensor, target_len: Optional[int] = None):
+        if target_len is None:
+            target_len = motion.shape[1]
+        latents = self.encode(motion)
+        return self.decode(latents, target_len), latents
+
+
+# =====================================================================
+# Quaternion helpers for joint recovery
+# =====================================================================
+
+def _qinv(q: torch.Tensor) -> torch.Tensor:
+    mask = torch.ones_like(q)
+    mask[..., 1:] = -mask[..., 1:]
+    return q * mask
+
+
+def _qrot(q: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    original_shape = list(v.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 recover_from_ric(features: torch.Tensor, num_joints: int) -> torch.Tensor:
+    """HumanML3D features -> joint positions.
+
+    `features` are the DENORMALISED 263-dimensional vectors; the return value
+    is `(..., time, num_joints, 3)` in metres.
+    """
+    rot_vel = features[..., 0]
+    r_rot_ang = torch.zeros_like(rot_vel)
+    r_rot_ang[..., 1:] = rot_vel[..., :-1]
+    r_rot_ang = torch.cumsum(r_rot_ang, dim=-1)
+
+    r_rot_quat = torch.zeros(features.shape[:-1] + (4,),
+                             dtype=features.dtype, device=features.device)
+    r_rot_quat[..., 0] = torch.cos(r_rot_ang)
+    r_rot_quat[..., 2] = torch.sin(r_rot_ang)
+
+    r_pos = torch.zeros(features.shape[:-1] + (3,),
+                        dtype=features.dtype, device=features.device)
+    r_pos[..., 1:, [0, 2]] = features[..., :-1, 1:3]
+    r_pos = _qrot(_qinv(r_rot_quat), r_pos)
+    r_pos = torch.cumsum(r_pos, dim=-2)
+    r_pos[..., 1] = features[..., 3]
+
+    positions = features[..., 4:(num_joints - 1) * 3 + 4]
+    positions = positions.view(positions.shape[:-1] + (-1, 3))
+    positions = _qrot(
+        _qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4,)), positions
+    )
+    positions[..., 0] += r_pos[..., 0:1]
+    positions[..., 2] += r_pos[..., 2:3]
+    return torch.cat([r_pos.unsqueeze(-2), positions], dim=-2)
+
+
+# =====================================================================
+# Model
+# =====================================================================
+
+class MugenOutput(ModelOutput):
+    """Output of :meth:`MugenForConditionalGeneration.forward`.
+
+    Attributes:
+        latents: `(batch, K, latent_dim)` slots, in raw autoencoder space.
+        reconstruction: `(batch, time, motion_input_dim)` when a motion was given.
+        logits: language-model logits when `input_ids` was given.
+    """
+
+    latents: Optional[torch.FloatTensor] = None
+    reconstruction: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+
+
+class MugenPreTrainedModel(PreTrainedModel):
+    config_class = MugenConfig
+    base_model_prefix = "mugen"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+
+class MugenForConditionalGeneration(MugenPreTrainedModel):
+    """MUGEN: text-to-motion generation and motion-to-text understanding.
+
+    The public API is three methods:
+
+    * :meth:`generate_motion` turns text into motion features.
+    * :meth:`generate_caption` turns motion into text.
+    * :meth:`encode_motion` / :meth:`decode_motion` expose the latent interface
+      both directions share.
+
+    :meth:`features_to_joints` converts motion features into joint positions
+    for rendering or measurement.
+    """
+
+    _tied_weights_keys = ["language_model.lm_head.weight"]
+    _keys_to_ignore_on_load_missing = [r"motion_autoencoder\.decoder_query_pe"]
+
+    def __init__(self, config: MugenConfig):
+        super().__init__(config)
+
+        self.k = int(config.k_latent_slots)
+        self.latent_dim = int(config.latent_dim)
+
+        # ---- part 1: the motion autoencoder ----
+        self.motion_autoencoder = AdaptiveLengthAutoEncoder(config)
+
+        # ---- part 2: the language model ----
+        gpt2_config = GPT2Config(
+            vocab_size=config.vocab_size,
+            n_positions=config.n_positions,
+            n_embd=config.n_embd,
+            n_layer=config.n_layer,
+            n_head=config.n_head,
+            n_inner=config.n_inner,
+            activation_function=config.activation_function,
+            resid_pdrop=config.resid_pdrop,
+            embd_pdrop=config.embd_pdrop,
+            attn_pdrop=config.attn_pdrop,
+            layer_norm_epsilon=config.layer_norm_epsilon,
+            initializer_range=config.initializer_range,
+            bos_token_id=config.bos_token_id,
+            eos_token_id=config.eos_token_id,
+            tie_word_embeddings=True,
+        )
+        self.language_model = GPT2LMHeadModel(gpt2_config)
+        self.n_layers = int(config.n_layer)
+        hidden_size = int(config.n_embd)
+
+        # ---- part 3a: the layer router ----
+        # A slot's routing logits are a bounded static table plus a bounded
+        # text-conditional correction. Both halves are tanh-capped, so no logit
+        # margin can grow large enough to saturate the routing softmax.
+        self.static_router_logits = nn.Parameter(torch.zeros(self.k, self.n_layers))
+        self.text_mem_proj = nn.Linear(hidden_size, self.latent_dim)
+        self.stage2_latent_queries = nn.Parameter(torch.zeros(self.k, self.latent_dim))
+        self.cross_attn_blocks = nn.ModuleList([
+            CrossAttentionBlock(
+                d_model=self.latent_dim,
+                nhead=config.alae_nhead,
+                dim_feedforward=config.alae_dim_feedforward,
+                dropout=config.alae_dropout,
+                activation=config.alae_activation,
+            )
+            for _ in range(int(config.num_cross_attn_layers))
+        ])
+        self.cross_attn_norm = nn.LayerNorm(self.latent_dim)
+        self.layer_router = nn.Sequential(
+            nn.Linear(self.latent_dim, int(config.router_hidden)),
+            nn.GELU(),
+            nn.Linear(int(config.router_hidden), self.n_layers),
+        )
+
+        # ---- part 3b: the calibrated latent head ----
+        self.projector_norm = nn.LayerNorm(hidden_size)
+        self.projector = nn.Linear(hidden_size, 2 * self.latent_dim)
+        self.feedback_norm = nn.LayerNorm(hidden_size)
+        self.latent_low_rank = int(config.latent_low_rank)
+        if self.latent_low_rank > 0:
+            self.latent_factors = nn.Parameter(
+                torch.zeros(self.k * self.latent_dim, self.latent_low_rank)
+            )
+            self.factor_scale_head = nn.Linear(hidden_size, self.latent_low_rank)
+
+        # ---- the understanding direction ----
+        self.motion_in_projector = nn.Linear(self.latent_dim, hidden_size)
+
+        # ---- normalisation statistics ----
+        # latent_*: standardise the latent space the head predicts in.
+        # feature_*: HumanML3D feature statistics, so the model can hand back
+        # denormalised motion without the caller owning the dataset.
+        self.register_buffer("latent_mean", torch.zeros(self.k, self.latent_dim))
+        self.register_buffer("latent_std", torch.ones(self.k, self.latent_dim))
+        self.register_buffer("feature_mean", torch.zeros(config.motion_input_dim))
+        self.register_buffer("feature_std", torch.ones(config.motion_input_dim))
+
+        self._tokenizer = None
+        self.post_init()
+
+    # ------------------------------------------------------------------
+    # plumbing
+    # ------------------------------------------------------------------
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def set_tokenizer(self, tokenizer):
+        """Cache a tokenizer so the generate helpers can be called without one."""
+        self._tokenizer = tokenizer
+        return self
+
+    def _get_tokenizer(self, tokenizer=None):
+        if tokenizer is not None:
+            return tokenizer
+        if self._tokenizer is None:
+            from transformers import AutoTokenizer
+            name = self.config._name_or_path or "zy22b/MUGEN"
+            self._tokenizer = AutoTokenizer.from_pretrained(name)
+        return self._tokenizer
+
+    # ------------------------------------------------------------------
+    # the latent interface both directions share
+    # ------------------------------------------------------------------
+    @torch.no_grad()
+    def encode_motion(self, motion: torch.Tensor, normalized: bool = False) -> torch.Tensor:
+        """Motion -> K latent slots.
+
+        Args:
+            motion: `(batch, time, 263)` HumanML3D features.
+            normalized: `True` if `motion` is already standardised by the
+                dataset statistics. `False` (the default) normalises it here.
+
+        Returns:
+            `(batch, K, latent_dim)` slots in raw autoencoder space.
+        """
+        motion = motion.to(self.device, dtype=torch.float32)
+        if not normalized:
+            motion = (motion - self.feature_mean) / self.feature_std
+        return self.motion_autoencoder.encode(motion)
+
+    @torch.no_grad()
+    def decode_motion(self, latents: torch.Tensor, length: int,
+                      denormalize: bool = True) -> torch.Tensor:
+        """K latent slots -> `(batch, length, 263)` motion features."""
+        latents = latents.to(self.device, dtype=torch.float32)
+        feats = self.motion_autoencoder.decode(latents, target_len=int(length))
+        if denormalize:
+            feats = feats * self.feature_std + self.feature_mean
+        return feats
+
+    def standardize_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        return (latents - self.latent_mean) / self.latent_std
+
+    def destandardize_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        return latents * self.latent_std + self.latent_mean
+
+    # ------------------------------------------------------------------
+    # text -> latents
+    # ------------------------------------------------------------------
+    def _build_t2m_prompts(self, texts: Sequence[str]) -> List[str]:
+        tpl = self.config.t2m_prompt_template
+        return [tpl.format(text=t.strip(), mot=self.config.mot_token) for t in texts]
+
+    def _routed_hidden(self, input_ids: torch.Tensor,
+                       attention_mask: torch.Tensor) -> torch.Tensor:
+        """Roll out K steps and return the routed per-slot hidden states.
+
+        Returns `(batch, K, hidden_size)`. Each slot is a soft mixture over all
+        `n_layers` transformer layers of its own rollout step, with mixture
+        weights produced by the router from the prompt.
+        """
+        K = self.k
+        device = input_ids.device
+        prompt_len = input_ids.shape[1]
+
+        def _positions(start: int, count: int) -> torch.Tensor:
+            return torch.arange(start, start + count, dtype=torch.long,
+                                device=device).unsqueeze(0)
+
+        out = self.language_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=_positions(0, prompt_len),
+            output_hidden_states=True,
+            use_cache=True,
+            return_dict=True,
+        )
+
+        # Router memory is the FULL last-layer token sequence plus its pad mask,
+        # so the routing queries can attend to individual words.
+        memory = self.text_mem_proj(out.hidden_states[-1])
+        pad_mask = attention_mask == 0
+
+        def _stack_all_layers(hidden_states) -> torch.Tensor:
+            cols = [h[:, -1:, :] for h in hidden_states[1:]]
+            return torch.stack(cols, dim=2)          # (B, 1, L, H)
+
+        all_steps = [_stack_all_layers(out.hidden_states)]
+        last_hiddens = [out.hidden_states[-1][:, -1:, :]]
+        past = out.past_key_values
+        cur_attn = attention_mask
+
+        for step_idx in range(K - 1):
+            next_input = self.feedback_norm(last_hiddens[-1])
+            cur_attn = torch.cat(
+                [cur_attn, torch.ones((cur_attn.shape[0], 1), dtype=cur_attn.dtype, device=device)],
+                dim=1,
+            )
+            step = self.language_model(
+                inputs_embeds=next_input,
+                attention_mask=cur_attn,
+                position_ids=_positions(prompt_len + step_idx, 1),
+                past_key_values=past,
+                output_hidden_states=True,
+                use_cache=True,
+                return_dict=True,
+            )
+            past = step.past_key_values
+            all_steps.append(_stack_all_layers(step.hidden_states))
+            last_hiddens.append(step.hidden_states[-1][:, -1:, :])
+
+        H_all = torch.cat(all_steps, dim=1)           # (B, K, L, H)
+        B = H_all.shape[0]
+
+        latents = self.stage2_latent_queries.unsqueeze(0).expand(B, -1, -1)
+        for block in self.cross_attn_blocks:
+            latents = block(latents, memory, memory_key_padding_mask=pad_mask)
+        latents = self.cross_attn_norm(latents)       # (B, K, D)
+
+        s_s = float(self.config.router_static_scale)
+        s_d = float(self.config.router_delta_scale)
+        static = s_s * torch.tanh(self.static_router_logits / s_s)      # (K, L)
+        delta = s_d * torch.tanh(self.layer_router(latents).float() / s_d)
+        logits = static.unsqueeze(0).float() + delta                    # (B, K, L)
+
+        tau = float(self.config.router_eval_tau)
+        weights = F.softmax(logits / tau, dim=-1).to(H_all.dtype)
+        return torch.einsum("bkl,bklh->bkh", weights, H_all)
+
+    def _latent_distribution(self, hidden: torch.Tensor):
+        """Routed hidden states -> `(mu, logvar, log_a)` in standardised space."""
+        h = self.projector_norm(hidden)
+        mu, logvar = self.projector(h).chunk(2, dim=-1)
+        logvar = logvar.clamp(-10.0, 10.0)
+        log_a = None
+        if self.latent_low_rank > 0:
+            log_a = self.factor_scale_head(h.mean(dim=1))               # (B, r)
+        return mu, logvar, log_a
+
+    def sample_latents(self, mu: torch.Tensor, logvar: torch.Tensor,
+                       log_a: Optional[torch.Tensor], temperature: float,
+                       generator: Optional[torch.Generator] = None) -> torch.Tensor:
+        """One structured draw `z = mu + t * (U (a * eps1) + sigma * eps2)`.
+
+        The rank-r factor term carries the variance that is correlated across
+        slots; the diagonal term carries the rest. `temperature` scales the
+        whole zero-mean perturbation, so 0.0 returns `mu` exactly.
+        """
+        B, K, D = mu.shape
+        if temperature <= 0.0:
+            return mu
+
+        def _randn(like):
+            if generator is None:
+                return torch.randn_like(like)
+            return torch.randn(like.shape, generator=generator,
+                               device=like.device, dtype=like.dtype)
+
+        # The factor noise is drawn BEFORE the diagonal noise, matching the
+        # training implementation, so a given seed reproduces the same draw
+        # here as it does there.
+        factor_part = 0.0
+        if log_a is not None:
+            a = torch.exp(log_a.clamp(-6.0, 4.0))                       # (B, r)
+            factor_part = torch.einsum(
+                "dr,br->bd", self.latent_factors, a * _randn(a)
+            )
+        sigma = torch.exp(0.5 * logvar).reshape(B, K * D)
+        diag_part = sigma * _randn(sigma)
+        z = mu.reshape(B, K * D) + float(temperature) * (factor_part + diag_part)
+        return z.reshape(B, K, D)
+
+    @torch.no_grad()
+    def text_to_latents(self, texts: Union[str, Sequence[str]], tokenizer=None,
+                        temperature: Optional[float] = None,
+                        generator: Optional[torch.Generator] = None) -> torch.Tensor:
+        """Text -> K latent slots in raw autoencoder space."""
+        if isinstance(texts, str):
+            texts = [texts]
+        tokenizer = self._get_tokenizer(tokenizer)
+        prev_side = tokenizer.padding_side
+        tokenizer.padding_side = "left"
+        try:
+            enc = tokenizer(
+                self._build_t2m_prompts(texts),
+                return_tensors="pt", padding=True, truncation=True, max_length=256,
+            )
+        finally:
+            tokenizer.padding_side = prev_side
+        enc = {k: v.to(self.device) for k, v in enc.items()}
+
+        hidden = self._routed_hidden(enc["input_ids"], enc["attention_mask"])
+        mu, logvar, log_a = self._latent_distribution(hidden)
+        temp = float(self.config.eval_sample_temperature if temperature is None else temperature)
+        z_norm = self.sample_latents(mu, logvar, log_a, temp, generator=generator)
+        return self.destandardize_latents(z_norm)
+
+    @torch.no_grad()
+    def generate_motion(self, texts: Union[str, Sequence[str]],
+                        lengths: Union[int, Sequence[int], None] = None,
+                        tokenizer=None,
+                        temperature: Optional[float] = None,
+                        denormalize: bool = True,
+                        generator: Optional[torch.Generator] = None) -> torch.Tensor:
+        """Text -> motion features.
+
+        Args:
+            texts: one description, or a batch of them.
+            lengths: target frame count per description (20 fps). A single int
+                applies to the whole batch. Defaults to 120 frames (6 seconds).
+                HumanML3D clips are multiples of 4 frames and at most 196, so
+                stay in that range to stay in distribution.
+            tokenizer: the tokenizer shipped with this repo. Loaded lazily when
+                omitted.
+            temperature: multiplier on the sampled perturbation. Defaults to
+                `config.eval_sample_temperature` (1.0), the calibrated
+                conditional distribution. 0.0 decodes the distribution mean,
+                which is not the protocol the published numbers use.
+            denormalize: return features in raw HumanML3D units. Keep this on
+                unless you intend to feed the output back into the encoder.
+            generator: optional `torch.Generator` for reproducible draws.
+
+        Returns:
+            `(batch, max_length, 263)`. When lengths differ within a batch, the
+            tensor is padded with zeros and each row is valid up to its own
+            length.
+        """
+        if isinstance(texts, str):
+            texts = [texts]
+        texts = list(texts)
+        if lengths is None:
+            lengths = [120] * len(texts)
+        elif isinstance(lengths, int):
+            lengths = [lengths] * len(texts)
+        lengths = [int(v) for v in lengths]
+        if len(lengths) != len(texts):
+            raise ValueError(f"got {len(texts)} texts but {len(lengths)} lengths")
+
+        latents = self.text_to_latents(texts, tokenizer=tokenizer,
+                                       temperature=temperature, generator=generator)
+
+        max_len = max(lengths)
+        if len(set(lengths)) == 1:
+            return self.decode_motion(latents, max_len, denormalize=denormalize)
+
+        # Mixed lengths: the decoder's relative-phase queries mean a row must be
+        # decoded at its own length, not cropped from the longest one.
+        out = latents.new_zeros((len(texts), max_len, self.config.motion_input_dim))
+        for i, length in enumerate(lengths):
+            out[i, :length] = self.decode_motion(
+                latents[i: i + 1], length, denormalize=denormalize
+            )[0]
+        return out
+
+    # ------------------------------------------------------------------
+    # motion -> text
+    # ------------------------------------------------------------------
+    def _m2t_prompt_embeds(self, latents: torch.Tensor, tokenizer):
+        """`[instruction, <MOT>, K motion slots]` as input embeddings."""
+        B = latents.shape[0]
+        device = latents.device
+        wte = self.language_model.get_input_embeddings()
+
+        prefix_ids = torch.tensor(
+            tokenizer(self.config.m2t_prompt_prefix, add_special_tokens=False)["input_ids"],
+            dtype=torch.long, device=device,
+        ).unsqueeze(0).expand(B, -1)
+        mot_ids = torch.full((B, 1), int(self.config.mot_token_id),
+                             dtype=torch.long, device=device)
+
+        prompt = torch.cat(
+            [wte(prefix_ids), wte(mot_ids), self.motion_in_projector(latents)], dim=1
+        )
+        attn = torch.ones(prompt.shape[:2], dtype=torch.long, device=device)
+        return prompt, attn
+
+    @torch.no_grad()
+    def generate_caption(self, motion: Optional[torch.Tensor] = None,
+                         latents: Optional[torch.Tensor] = None,
+                         tokenizer=None, normalized: bool = False,
+                         max_new_tokens: Optional[int] = None,
+                         num_beams: Optional[int] = None) -> List[str]:
+        """Motion -> one caption per clip.
+
+        Pass either `motion` (`(batch, time, 263)` features) or `latents`
+        (`(batch, K, latent_dim)`) if you already have the slots. Set
+        `normalized=True` when `motion` is already standardised.
+        """
+        if (motion is None) == (latents is None):
+            raise ValueError("pass exactly one of `motion` or `latents`")
+        tokenizer = self._get_tokenizer(tokenizer)
+        if latents is None:
+            latents = self.encode_motion(motion, normalized=normalized)
+        latents = latents.to(self.device, dtype=torch.float32)
+
+        prompt, attn = self._m2t_prompt_embeds(latents, tokenizer)
+        gen = self.language_model.generate(
+            inputs_embeds=prompt,
+            attention_mask=attn,
+            max_new_tokens=int(self.config.m2t_max_new_tokens
+                               if max_new_tokens is None else max_new_tokens),
+            do_sample=False,
+            num_beams=int(self.config.m2t_num_beams if num_beams is None else num_beams),
+            pad_token_id=tokenizer.eos_token_id,
+        )
+        # Called with inputs_embeds and no input_ids, generate returns only the
+        # newly produced tokens.
+        return [t.strip() for t in tokenizer.batch_decode(gen, skip_special_tokens=True)]
+
+    # ------------------------------------------------------------------
+    # utilities
+    # ------------------------------------------------------------------
+    def features_to_joints(self, features: torch.Tensor,
+                          normalized: bool = False) -> torch.Tensor:
+        """Motion features -> `(batch, time, num_joints, 3)` joint positions in metres."""
+        features = features.to(self.device, dtype=torch.float32)
+        if normalized:
+            features = features * self.feature_std + self.feature_mean
+        return recover_from_ric(features, int(self.config.num_joints))
+
+    def forward(self,
+                motion: Optional[torch.Tensor] = None,
+                input_ids: Optional[torch.LongTensor] = None,
+                attention_mask: Optional[torch.Tensor] = None,
+                target_length: Optional[int] = None,
+                normalized: bool = False,
+                return_dict: bool = True,
+                **kwargs) -> MugenOutput:
+        """Autoencode a motion, run the language model, or both.
+
+        This is the plumbing entry point. For generation use
+        :meth:`generate_motion` and :meth:`generate_caption`.
+        """
+        latents = reconstruction = logits = None
+        if motion is not None:
+            motion = motion.to(self.device, dtype=torch.float32)
+            if not normalized:
+                motion = (motion - self.feature_mean) / self.feature_std
+            latents = self.motion_autoencoder.encode(motion)
+            reconstruction = self.motion_autoencoder.decode(
+                latents, target_len=int(target_length or motion.shape[1])
+            )
+        if input_ids is not None:
+            logits = self.language_model(
+                input_ids=input_ids, attention_mask=attention_mask, return_dict=True
+            ).logits
+        return MugenOutput(latents=latents, reconstruction=reconstruction, logits=logits)
+
+
+MugenConfig.register_for_auto_class()
+MugenForConditionalGeneration.register_for_auto_class("AutoModelForCausalLM")

special_tokens_map.json

@@ -0,0 +1,15 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<MOT>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

tokenizer_config.json

@@ -0,0 +1,57 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50257": {
+      "content": "<motion_id_0>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50258": {
+      "content": "<motion_id_1>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50259": {
+      "content": "<motion_id_2>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50260": {
+      "content": "<MOT>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<MOT>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "legacy": true,
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}