Release TiTok motion encoder code and checkpoint

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [2023] [Zhongjie Duan]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,6 @@
+IM-Animation motion encoder uses modified TiTok components from
+https://github.com/bytedance/1d-tokenizer (Apache-2.0).
+Original copyright and reference notices are retained in source files.
+This distribution extracts the encoder, vector quantizer, and preprocessing
+from the IM-Animation research implementation and adds an encoder-only loader.
+The original TiTok legacy reshape is preserved.

README.md

@@ -0,0 +1,87 @@
+---
+language:
+- en
+license: apache-2.0
+tags:
+- im-animation
+- titok
+- motion-encoder
+- image-feature-extraction
+- arxiv:2602.07498
+---
+
+# IM-Animation
+
+[Paper](https://arxiv.org/abs/2602.07498) · [Project page](https://rabberk.github.io/IM-Animation/) · [Motion encoder weights](https://huggingface.co/Rbaerk/IM-Animation-Motion-Encoder)
+
+**IM-Animation: An Implicit Motion Representation for Identity-decoupled Character Animation**
+
+This release contains the final locally retained **TiTok-based motion encoder** implementation and an exported checkpoint. It provides frame-level motion tokens; it does not include the full animation generator or retargeting network. Code and project videos are available in the linked GitHub repository.
+
+## Installation
+
+```bash
+git clone https://github.com/rabberk/IM-Animation.git
+cd IM-Animation
+pip install -r requirements.txt
+```
+
+The encoder was verified with PyTorch 2.7.1 on CPU using the exported BF16 weights.
+
+## Download weights
+
+```python
+from huggingface_hub import hf_hub_download
+
+hf_hub_download(
+    repo_id="Rbaerk/IM-Animation-Motion-Encoder",
+    filename="motion_encoder_latest.safetensors",
+    local_dir=".",
+)
+```
+
+The weight file is 607,017,208 bytes (approximately 579 MiB). Its SHA256 and provenance are in [`checkpoint_info.json`](checkpoint_info.json).
+
+## Encode frames
+
+Run from the repository directory:
+
+```python
+import torch
+from motion_encoder import MotionEncoder
+
+model = MotionEncoder.from_pretrained(device="cpu")  # or device="cuda"
+frames = torch.rand(1, 3, 256, 256).to(
+    device=next(model.parameters()).device,
+    dtype=next(model.parameters()).dtype,
+)
+with torch.inference_mode():
+    tokens, metrics = model.encode(frames)
+print(tokens.shape)  # [1, 12, 1, 32]
+```
+
+`frames` must be RGB, normalized to `[0, 1]`, with shape `[N, 3, 256, 256]`.
+The training preprocessing pads portrait frames horizontally to a square and resizes them to 256×256 with bilinear interpolation (`align_corners=False`). The original `HW_encoder_2` preprocessing class is included in `encoder_blocks.py`.
+
+Each frame produces 32 tokens of 12 dimensions. The training integration flattens them into 384 dimensions per frame before retargeting. The encoder processes frames independently; temporal retargeting is outside this release.
+
+## Architecture
+
+- TiTokEncoder: 24 Transformer layers, hidden width 1024, 16 attention heads.
+- Patch size 16; 32 learned latent tokens.
+- 12-dimensional output projection and a 4096-entry vector-quantization codebook.
+- Original `is_legacy=True` token reshape is retained for checkpoint compatibility.
+
+Implementation: `motion_encoder.py` assembles the modules; `encoder_blocks.py` contains the original encoder and preprocessing; `quantizer.py` contains the original VQ implementation.
+
+## Checkpoint provenance and verification
+
+The selected checkpoint is `train_dit_5C_v6_part5/step-12200.safetensors`, dated 2025-11-28 UTC by file modification time. It is the newest checkpoint in the inspected local runs, rather than a claim of best quality or a verified paper-final checkpoint.
+
+Training saved only trainable parameters. The selected checkpoint supplies 300 encoder/latent-token tensors. The frozen VQ codebook is restored from `train_motion_only_full_3C_20joint/step-3700.safetensors`, following the available training initialization code. This yields a complete 301-tensor encoding module. The historical run's frozen state has not been independently verified.
+
+Validation checked tensor byte hashes against their source checkpoints, strict state-dict loading, and exact single-frame output agreement with the available original TiTok implementation. Full video-generation quality was not evaluated in this export.
+
+## Acknowledgments and license
+
+This encoder builds on [TiTok / 1d-tokenizer](https://github.com/bytedance/1d-tokenizer). Source attribution is preserved. See [LICENSE](LICENSE) and [NOTICE](NOTICE).

checkpoint_info.json

@@ -0,0 +1,24 @@
+{
+  "checkpoint": "train_dit_5C_v6_part5/step-12200.safetensors",
+  "checkpoint_mtime_utc": "2025-11-28T03:50:08Z",
+  "selection": "Latest file modification time among the inspected local checkpoints; not a quality ranking.",
+  "weight_file": "motion_encoder_latest.safetensors",
+  "weight_bytes": 607017208,
+  "sha256": "008316ed723d89a5cc4b5d6b2917509b67b434fdba92c722b36116bc0aa32c10",
+  "saved_tensors_from_selected_checkpoint": 300,
+  "frozen_codebook_source": "train_motion_only_full_3C_20joint/step-3700.safetensors",
+  "total_exported_tensors": 301,
+  "limitations": "Frozen VQ codebook reconstructed according to the available training script; historical runtime state has not been independently verified.",
+  "verification": {
+    "tensor_hashes_verified": 301,
+    "strict_load": true,
+    "original_forward_exact_match": true,
+    "output_shape": [
+      1,
+      12,
+      1,
+      32
+    ],
+    "torch": "2.7.1+cu118"
+  }
+}

config.yaml

@@ -0,0 +1,15 @@
+model:
+  vq_model:
+    codebook_size: 4096
+    token_size: 12
+    use_l2_norm: true
+    commitment_cost: 0.25
+    vit_enc_model_size: large
+    vit_enc_patch_size: 16
+    num_latent_tokens: 32
+    quantize_mode: vq
+    is_legacy: true
+dataset:
+  preprocessing:
+    height_size: 256
+    width_size: 256

encoder_blocks.py

@@ -0,0 +1,197 @@
+# Extracted for IM-Animation: encoder and preprocessing only; original forward logic retained.
+"""Building blocks for TiTok.
+
+Copyright (2024) Bytedance Ltd. and/or its affiliates
+
+Licensed under the Apache License, Version 2.0 (the "License"); 
+you may not use this file except in compliance with the License. 
+You may obtain a copy of the License at 
+
+    http://www.apache.org/licenses/LICENSE-2.0 
+
+Unless required by applicable law or agreed to in writing, software 
+distributed under the License is distributed on an "AS IS" BASIS, 
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+See the License for the specific language governing permissions and 
+limitations under the License. 
+
+Reference: 
+    https://github.com/mlfoundations/open_clip/blob/main/src/open_clip/transformer.py
+    https://github.com/baofff/U-ViT/blob/main/libs/timm.py
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from collections import OrderedDict
+from einops import rearrange
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+            self,
+            d_model,
+            n_head,
+            mlp_ratio = 4.0,
+            act_layer = nn.GELU,
+            norm_layer = nn.LayerNorm
+        ):
+        super().__init__()
+
+        self.ln_1 = norm_layer(d_model)
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.mlp_ratio = mlp_ratio
+        # optionally we can disable the FFN
+        if mlp_ratio > 0:
+            self.ln_2 = norm_layer(d_model)
+            mlp_width = int(d_model * mlp_ratio)
+            self.mlp = nn.Sequential(OrderedDict([
+                ("c_fc", nn.Linear(d_model, mlp_width)),
+                ("gelu", act_layer()),
+                ("c_proj", nn.Linear(mlp_width, d_model))
+            ]))
+
+    def attention(
+            self,
+            x: torch.Tensor
+    ):
+        return self.attn(x, x, x, need_weights=False)[0]
+
+    def forward(
+            self,
+            x: torch.Tensor,
+    ):
+        attn_output = self.attention(x=self.ln_1(x))
+        x = x + attn_output
+        if self.mlp_ratio > 0:
+            x = x + self.mlp(self.ln_2(x))
+        return x
+
+def _expand_token(token, batch_size: int):
+    return token.unsqueeze(0).expand(batch_size, -1, -1)
+
+class TiTokEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.width_size = config.dataset.preprocessing.width_size 
+        self.height_size = config.dataset.preprocessing.height_size 
+        self.patch_size = config.model.vq_model.vit_enc_patch_size
+        self.grid_size_w = self.width_size // self.patch_size 
+        self.grid_size_h = self.height_size // self.patch_size 
+        self.model_size = config.model.vq_model.vit_enc_model_size
+        self.num_latent_tokens = config.model.vq_model.num_latent_tokens
+        self.token_size = config.model.vq_model.token_size
+
+        if config.model.vq_model.get("quantize_mode", "vq") == "vae":
+            self.token_size = self.token_size * 2 # needs to split into mean and std
+
+        self.is_legacy = config.model.vq_model.get("is_legacy", True)
+
+        self.width = {
+                "small": 512,
+                "base": 768,
+                "large": 1024,
+            }[self.model_size]
+        self.num_layers = {
+                "small": 8,
+                "base": 12,
+                "large": 24,
+            }[self.model_size]
+        self.num_heads = {
+                "small": 8,
+                "base": 12,
+                "large": 16,
+            }[self.model_size]
+        
+        self.patch_embed = nn.Conv2d(
+            in_channels=3, out_channels=self.width,
+              kernel_size=self.patch_size, stride=self.patch_size,padding = (4,2), bias=True)
+        
+        scale = self.width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(1, self.width))
+        self.positional_embedding = nn.Parameter(
+                scale * torch.randn(self.grid_size_h*self.grid_size_w + 1, self.width))
+        self.latent_token_positional_embedding = nn.Parameter(
+            scale * torch.randn(self.num_latent_tokens, self.width))
+        self.ln_pre = nn.LayerNorm(self.width)
+        self.transformer = nn.ModuleList()
+        for i in range(self.num_layers):
+            self.transformer.append(ResidualAttentionBlock(
+                self.width, self.num_heads, mlp_ratio=4.0
+            ))
+        self.ln_post = nn.LayerNorm(self.width)
+        self.conv_out = nn.Conv2d(self.width, self.token_size, kernel_size=1, bias=True)
+
+    def forward(self, pixel_values, latent_tokens):
+        batch_size = pixel_values.shape[0]
+        x = pixel_values
+        x = self.patch_embed(x)
+        x = x.reshape(x.shape[0], x.shape[1], -1)
+        x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
+        # class embeddings and positional embeddings
+        x = torch.cat([_expand_token(self.class_embedding, x.shape[0]).to(x.dtype), x], dim=1)
+        x = x + self.positional_embedding.to(x.dtype) # shape = [*, grid ** 2 + 1, width]
+        
+
+        latent_tokens = _expand_token(latent_tokens, x.shape[0]).to(x.dtype)
+        latent_tokens = latent_tokens + self.latent_token_positional_embedding.to(x.dtype)
+        x = torch.cat([x, latent_tokens], dim=1)
+        def create_custom_forward(module):
+            def custom_forward(*inputs):
+                return module(*inputs)
+            return custom_forward
+        x = self.ln_pre(x)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        for i in range(self.num_layers):
+            # x = self.transformer[i](x)
+            #with torch.autograd.graph.save_on_cpu():
+                x = torch.utils.checkpoint.checkpoint(create_custom_forward(self.transformer[i]), x,use_reentrant=False)
+
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        
+        latent_tokens = x[:, 1+self.grid_size_h*self.grid_size_w:]
+        latent_tokens = self.ln_post(latent_tokens)
+        # fake 2D shape
+        if self.is_legacy:
+            latent_tokens = latent_tokens.reshape(batch_size, self.width, self.num_latent_tokens, 1)
+        else:
+            # Fix legacy problem.
+            latent_tokens = latent_tokens.reshape(batch_size, self.num_latent_tokens, self.width, 1).permute(0, 2, 1, 3)
+        latent_tokens = self.conv_out(latent_tokens)
+        latent_tokens = latent_tokens.reshape(batch_size, self.token_size, 1, self.num_latent_tokens)
+        return latent_tokens
+
+class HW_encoder_2(nn.Module):
+    def __init__(self, in_channels):
+        super(HW_encoder_2, self).__init__()
+    #     self.conv0 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
+
+    #     self.conv1 = nn.Conv2d(in_channels*4, in_channels, kernel_size=3, padding=1)
+    #     self.conv2 = nn.Conv2d(in_channels*4 , in_channels, kernel_size=3, padding=1)
+    #     self.conv3 = nn.Conv2d(in_channels*4 , in_channels , kernel_size=3, padding=1)
+    #     # self.conv4 = nn.Conv2d(in_channels , in_channels//4 , kernel_size=3, padding=1)
+
+    # def pixel_shuffle(self, x, scale_factor=0.5):
+    #     n, c, h, w = x.size() 
+    #     new_h = int(h * scale_factor)
+    #     new_w = int(w * scale_factor)
+    #     x = x.view(n, int(c / (scale_factor ** 2)), new_h, new_w)
+
+        # return x
+    def forward(self, x):
+        B, C, T, H, W = x.shape
+        x = rearrange(x, "b c f h w -> (b f) c h w")
+
+        # Step 1: Pad the width from 480 to 832
+        padding_width = (H - W) // 2
+        x = F.pad(x, (padding_width, padding_width, 0, 0))  # Pad width only
+
+        # Step 2: Resize to target width 256
+        target_size = (256, 256)
+        x = F.interpolate(x, size=target_size, mode='bilinear', align_corners=False)
+
+        # Rearrange back to original shape
+        x = rearrange(x, "(b f) c h w -> b c f h w", f=T)
+
+        return x

motion_encoder.py

@@ -0,0 +1,38 @@
+"""Encoder-only assembly of the original TiTok motion encoding path.
+
+TiTokEncoder, VectorQuantizer, HW_encoder_2 retain the original implementations.
+Input to encode: preprocessed frames [N, 3, 256, 256], RGB in [0, 1].
+"""
+from pathlib import Path
+import torch
+from torch import nn
+from omegaconf import OmegaConf
+from safetensors.torch import load_file
+from encoder_blocks import TiTokEncoder, HW_encoder_2
+from quantizer import VectorQuantizer
+
+class MotionEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.encoder = TiTokEncoder(config)
+        vq = config.model.vq_model
+        self.latent_tokens = nn.Parameter(torch.empty(vq.num_latent_tokens, self.encoder.width))
+        self.quantize = VectorQuantizer(codebook_size=vq.codebook_size,
+            token_size=vq.token_size, commitment_cost=vq.commitment_cost,
+            use_l2_norm=vq.use_l2_norm)
+
+    def encode(self, frames):
+        z = self.encoder(pixel_values=frames, latent_tokens=self.latent_tokens)
+        return self.quantize(z)
+
+    def forward(self, frames):
+        return self.encode(frames)
+
+    @classmethod
+    def from_pretrained(cls, directory=None, device='cpu'):
+        directory = Path(directory or Path(__file__).parent)
+        config = OmegaConf.load(directory / 'config.yaml')
+        with torch.device('meta'):
+            model = cls(config)
+        model.load_state_dict(load_file(str(directory / 'motion_encoder_latest.safetensors')), strict=True, assign=True)
+        return model.to(device).eval()

motion_encoder_latest.safetensors

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

quantizer.py

@@ -0,0 +1,170 @@
+"""Vector quantizer.
+
+Copyright (2024) Bytedance Ltd. and/or its affiliates
+
+Licensed under the Apache License, Version 2.0 (the "License"); 
+you may not use this file except in compliance with the License. 
+You may obtain a copy of the License at 
+
+    http://www.apache.org/licenses/LICENSE-2.0 
+
+Unless required by applicable law or agreed to in writing, software 
+distributed under the License is distributed on an "AS IS" BASIS, 
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+See the License for the specific language governing permissions and 
+limitations under the License.
+
+Reference: 
+    https://github.com/CompVis/taming-transformers/blob/master/taming/modules/vqvae/quantize.py
+    https://github.com/google-research/magvit/blob/main/videogvt/models/vqvae.py
+    https://github.com/CompVis/latent-diffusion/blob/main/ldm/modules/distributions/distributions.py
+    https://github.com/lyndonzheng/CVQ-VAE/blob/main/quantise.py
+"""
+from typing import Mapping, Text, Tuple
+
+import torch
+from einops import rearrange
+from accelerate.utils.operations import gather
+from torch.cuda.amp import autocast
+
+class VectorQuantizer(torch.nn.Module):
+    def __init__(self,
+                 codebook_size: int = 1024,
+                 token_size: int = 256,
+                 commitment_cost: float = 0.25,
+                 use_l2_norm: bool = False,
+                 clustering_vq: bool = False
+                 ):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.token_size = token_size
+        self.commitment_cost = commitment_cost
+
+        self.embedding = torch.nn.Embedding(codebook_size, token_size)
+        self.embedding.weight.data.uniform_(-1.0 / codebook_size, 1.0 / codebook_size)
+        self.use_l2_norm = use_l2_norm
+
+        self.clustering_vq = clustering_vq
+        if clustering_vq:
+            self.decay = 0.99
+            self.register_buffer("embed_prob", torch.zeros(self.codebook_size))
+
+    # Ensure quantization is performed using f32
+    # @autocast(enabled=False)
+    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, Mapping[Text, torch.Tensor]]:
+        # z = z.float()
+        z = rearrange(z, 'b c h w -> b h w c').contiguous()
+        z_flattened = rearrange(z, 'b h w c -> (b h w) c')
+        unnormed_z_flattened = z_flattened
+
+        if self.use_l2_norm:
+            z_flattened = torch.nn.functional.normalize(z_flattened, dim=-1)
+            embedding = torch.nn.functional.normalize(self.embedding.weight, dim=-1)
+        else:
+            embedding = self.embedding.weight
+        d = torch.sum(z_flattened**2, dim=1, keepdim=True) + \
+            torch.sum(embedding**2, dim=1) - 2 * \
+            torch.einsum('bd,dn->bn', z_flattened, embedding.T)
+
+        min_encoding_indices = torch.argmin(d, dim=1) # num_ele
+        z_quantized = self.get_codebook_entry(min_encoding_indices).view(z.shape)
+
+        if self.use_l2_norm:
+            z = torch.nn.functional.normalize(z, dim=-1)
+
+        # compute loss for embedding
+        commitment_loss = self.commitment_cost * torch.mean((z_quantized.detach() - z) **2)
+        codebook_loss = torch.mean((z_quantized - z.detach()) **2)
+
+        if self.clustering_vq and self.training:
+            with torch.no_grad():
+                # Gather distance matrix from all GPUs.
+                encoding_indices = gather(min_encoding_indices)
+                if len(min_encoding_indices.shape) != 1:
+                    raise ValueError(f"min_encoding_indices in a wrong shape, {min_encoding_indices.shape}")
+                # Compute and update the usage of each entry in the codebook.
+                encodings = torch.zeros(encoding_indices.shape[0], self.codebook_size, device=z.device)
+                encodings.scatter_(1, encoding_indices.unsqueeze(1), 1)
+                avg_probs = torch.mean(encodings, dim=0)
+                self.embed_prob.mul_(self.decay).add_(avg_probs, alpha=1-self.decay)
+                # Closest sampling to update the codebook.
+                all_d = gather(d)
+                all_unnormed_z_flattened = gather(unnormed_z_flattened).detach()
+                if all_d.shape[0] != all_unnormed_z_flattened.shape[0]:
+                    raise ValueError(
+                        "all_d and all_unnormed_z_flattened have different length" + 
+                        f"{all_d.shape}, {all_unnormed_z_flattened.shape}")
+                indices = torch.argmin(all_d, dim=0)
+                random_feat = all_unnormed_z_flattened[indices]
+                # Decay parameter based on the average usage.
+                decay = torch.exp(-(self.embed_prob * self.codebook_size * 10) /
+                                   (1 - self.decay) - 1e-3).unsqueeze(1).repeat(1, self.token_size)
+                self.embedding.weight.data = self.embedding.weight.data * (1 - decay) + random_feat * decay
+
+        loss = commitment_loss + codebook_loss
+
+        # preserve gradients
+        z_quantized = z + (z_quantized - z).detach()
+
+        # reshape back to match original input shape
+        z_quantized = rearrange(z_quantized, 'b h w c -> b c h w').contiguous()
+
+        result_dict = dict(
+            quantizer_loss=loss,
+            commitment_loss=commitment_loss,
+            codebook_loss=codebook_loss,
+            min_encoding_indices=min_encoding_indices.view(z_quantized.shape[0], z_quantized.shape[2], z_quantized.shape[3])
+        )
+
+        return z_quantized, result_dict
+
+    def get_codebook_entry(self, indices):
+        if len(indices.shape) == 1:
+            z_quantized = self.embedding(indices)
+        elif len(indices.shape) == 2:
+            z_quantized = torch.einsum('bd,dn->bn', indices, self.embedding.weight)
+        else:
+            raise NotImplementedError
+        if self.use_l2_norm:
+            z_quantized = torch.nn.functional.normalize(z_quantized, dim=-1)
+        return z_quantized
+    
+
+class DiagonalGaussianDistribution(object):
+    @autocast(enabled=False)
+    def __init__(self, parameters, deterministic=False):
+        """Initializes a Gaussian distribution instance given the parameters.
+
+        Args:
+            parameters (torch.Tensor): The parameters for the Gaussian distribution. It is expected
+                to be in shape [B, 2 * C, *], where B is batch size, and C is the embedding dimension.
+                First C channels are used for mean and last C are used for logvar in the Gaussian distribution.
+            deterministic (bool): Whether to use deterministic sampling. When it is true, the sampling results
+                is purely based on mean (i.e., std = 0).
+        """
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters.float(), 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+    @autocast(enabled=False)
+    def sample(self):
+        x = self.mean.float() + self.std.float() * torch.randn(self.mean.shape).to(device=self.parameters.device)
+        return x
+
+    @autocast(enabled=False)
+    def mode(self):
+        return self.mean
+
+    @autocast(enabled=False)
+    def kl(self):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        else:
+            return 0.5 * torch.sum(torch.pow(self.mean.float(), 2)
+                                    + self.var.float() - 1.0 - self.logvar.float(),
+                                    dim=[1, 2])

requirements.txt

@@ -0,0 +1,6 @@
+torch>=2.1
+einops>=0.7
+omegaconf>=2.3
+safetensors>=0.4
+accelerate>=0.26
+huggingface_hub>=0.24