Wojtekb30 commited on
Commit
811f3dc
·
1 Parent(s): 829b0a7

Added new files

Browse files
.gitattributes CHANGED
@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
37
+ generated_motion.gif filter=lfs diff=lfs merge=lfs -text
38
+ walker.jpeg filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1 @@
 
 
1
+ .ipynb_checkpoints
RunVLA.py ADDED
@@ -0,0 +1,219 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import re
4
+ import matplotlib.pyplot as plt
5
+ import matplotlib.animation as animation
6
+ import torch.nn.functional as F
7
+ from PIL import Image # Used for loading image files
8
+ from IPython.display import display, Image as IPyImage # Used for displaying GIFs inside Jupyter notebooks
9
+
10
+ # Import classes required for the Vision-Language model
11
+ from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
12
+ from safetensors.torch import load_file
13
+ from rvq_model.rvq_model import MotionRVQ_VAE
14
+
15
+ # ==========================================
16
+ # 1. Configuration
17
+ # ==========================================
18
+ IMAGE_PATH = "walker.jpeg"
19
+ PROMPT_TEXT = "Describe the image and generate a matching physical motion."
20
+ QWEN_PATH = "./"
21
+ RVQ_WEIGHTS = "./rvq_model/motion_rvq_weights.safetensors"
22
+ OUTPUT_GIF = "generated_motion.gif"
23
+
24
+ torch.manual_seed(42)
25
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
26
+
27
+ # ==========================================
28
+ # 2. Load Qwen-VL (Brain & Eyes)
29
+ # ==========================================
30
+ print("Loading Qwen2.5-VL model and processor...")
31
+
32
+ # The processor includes both the tokenizer and image preprocessing pipeline
33
+ processor = AutoProcessor.from_pretrained(QWEN_PATH)
34
+
35
+ llm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
36
+ QWEN_PATH,
37
+ torch_dtype=torch.bfloat16,
38
+ device_map="auto" # Automatically places the model on the GPU if available
39
+ )
40
+ llm_model.eval()
41
+
42
+ # ==========================================
43
+ # 3. Load the RVQ model (Body)
44
+ # ==========================================
45
+ print("Loading RVQ decoder...")
46
+
47
+ rvq_model = MotionRVQ_VAE().to(device)
48
+ state_dict = load_file(RVQ_WEIGHTS, device=str(device))
49
+ rvq_model.load_state_dict(state_dict)
50
+ rvq_model.eval()
51
+
52
+ mean = np.load('./rvq_model/Mean.npy')
53
+ std = np.load('./rvq_model/Std.npy')
54
+
55
+ # ==========================================
56
+ # 4. Process image and generate a response
57
+ # ==========================================
58
+ print(f"\nAI instruction: '{PROMPT_TEXT}' with image '{IMAGE_PATH}'")
59
+
60
+ # Load the image using PIL
61
+ image = Image.open(IMAGE_PATH).convert("RGB")
62
+
63
+ # Vision-language prompt structure
64
+ messages = [
65
+ {
66
+ "role": "system",
67
+ "content": "You are an embodied AI. You reason about your physical state and output precise motor actions inside <move></move> tags."
68
+ },
69
+ {
70
+ "role": "user",
71
+ "content": [
72
+ {"type": "image"}, # Placeholder indicating where the image will be inserted
73
+ {"type": "text", "text": PROMPT_TEXT}
74
+ ]
75
+ }
76
+ ]
77
+
78
+ # Prepare both text and image inputs using the processor
79
+ text_prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
80
+ inputs = processor(
81
+ text=[text_prompt],
82
+ images=[image],
83
+ padding=True,
84
+ return_tensors="pt"
85
+ ).to(device)
86
+
87
+ print("Observing image, reasoning, and generating motion...")
88
+
89
+ with torch.no_grad():
90
+ outputs = llm_model.generate(
91
+ **inputs,
92
+ max_new_tokens=1024,
93
+ temperature=0.5, # Lower temperature improves RVQ token consistency
94
+ do_sample=True,
95
+ )
96
+
97
+ # Extract only the generated portion, excluding the original prompt
98
+ generated_ids = [
99
+ output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, outputs)
100
+ ]
101
+ response = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
102
+ response = response.replace("<|im_end|>", "").strip()
103
+
104
+ print("\n=== QWEN RESPONSE ===")
105
+ print(response)
106
+ print("=====================\n")
107
+
108
+ # ==========================================
109
+ # 5. Extract and parse motion tokens
110
+ # ==========================================
111
+ move_blocks = re.findall(r'<move>(.*?)</move>', response, re.DOTALL)
112
+
113
+ if not move_blocks:
114
+ print("ERROR: Qwen did not generate any motion tokens.")
115
+ exit()
116
+
117
+ all_tokens = []
118
+ for block in move_blocks:
119
+ tokens = re.findall(r'<m_(\d+)_(\d+)>', block)
120
+ all_tokens.extend(tokens)
121
+
122
+ if not all_tokens:
123
+ print("ERROR: No valid tokens were found inside the <move> tags.")
124
+ exit()
125
+
126
+ num_frames = len(all_tokens) // 4
127
+ token_matrix = np.zeros((4, num_frames), dtype=np.int64)
128
+
129
+ for i in range(num_frames):
130
+ for lvl in range(4):
131
+ token_idx = i * 4 + lvl
132
+ if token_idx < len(all_tokens):
133
+ parsed_lvl, val = all_tokens[token_idx]
134
+ token_matrix[lvl, i] = int(val)
135
+
136
+ token_tensor = torch.tensor(token_matrix, device=device).unsqueeze(0)
137
+
138
+ # ==========================================
139
+ # 6. Decode tokens into 3D motion
140
+ # ==========================================
141
+ print(f"Decoding {num_frames} token frames into 3D motion...")
142
+
143
+ with torch.no_grad():
144
+ z_q = 0
145
+ for lvl in range(4):
146
+ indices = token_tensor[:, lvl, :]
147
+ quantizer = rvq_model.rvq.quantizers[lvl]
148
+ level_z_q = F.embedding(indices, quantizer.embedding)
149
+ level_z_q = level_z_q.permute(0, 2, 1)
150
+ z_q = z_q + level_z_q
151
+ reconstructed_motion = rvq_model.decoder(z_q)
152
+
153
+ recon_data = reconstructed_motion.squeeze(0).permute(1, 0).cpu().numpy()
154
+ recon_data = (recon_data * std) + mean
155
+ T_frames = recon_data.shape[0]
156
+
157
+ # ==========================================
158
+ # 7. Save and visualize the generated motion
159
+ # ==========================================
160
+ def get_3d_joints(data_263):
161
+ frames = data_263.shape[0]
162
+ joints = np.zeros((frames, 22, 3))
163
+ for i in range(frames):
164
+ root_y = data_263[i, 3]
165
+ joints[i, 0] = [0, root_y, 0]
166
+ local_positions = data_263[i, 4:67].reshape(21, 3)
167
+ joints[i, 1:] = local_positions + [0, root_y, 0]
168
+ return joints
169
+
170
+ joints_recon = get_3d_joints(recon_data)
171
+
172
+ kinematic_tree = [
173
+ [0, 1, 4, 7, 10],
174
+ [0, 2, 5, 8, 11],
175
+ [0, 3, 6, 9, 12, 15],
176
+ [9, 13, 16, 18, 20],
177
+ [9, 14, 17, 19, 21]
178
+ ]
179
+
180
+ fig = plt.figure(figsize=(6, 6))
181
+ ax = fig.add_subplot(111, projection='3d')
182
+
183
+ def update(frame):
184
+ ax.clear()
185
+ ax.set_title(f"QWEN-VLA GENERATED MOTION\nFrame: {frame}/{T_frames}")
186
+ ax.set_xlim(-1, 1)
187
+ ax.set_ylim(-1, 1)
188
+ ax.set_zlim(0, 2)
189
+ ax.view_init(elev=10., azim=-90)
190
+ ax.axis('off')
191
+
192
+ for chain in kinematic_tree:
193
+ ax.plot(
194
+ joints_recon[frame, chain, 0],
195
+ joints_recon[frame, chain, 2],
196
+ joints_recon[frame, chain, 1],
197
+ linewidth=3,
198
+ marker='o',
199
+ markersize=4,
200
+ color='red'
201
+ )
202
+
203
+ print("Generating GIF file...")
204
+
205
+ ani = animation.FuncAnimation(
206
+ fig,
207
+ update,
208
+ frames=T_frames,
209
+ interval=50,
210
+ repeat=True
211
+ )
212
+
213
+ # Save using PillowWriter (works well in Jupyter and does not require FFmpeg)
214
+ ani.save(OUTPUT_GIF, writer='pillow', fps=20)
215
+
216
+ # Close the figure to prevent an empty plot from appearing in notebook output
217
+ plt.close(fig)
218
+
219
+ print(f"Animation saved to file: {OUTPUT_GIF}")
TrainVLA.py ADDED
@@ -0,0 +1,103 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from unsloth import FastVisionModel
2
+ import torch
3
+ from datasets import load_dataset
4
+ from transformers import TrainingArguments, TrainerCallback
5
+ from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
6
+
7
+ # Callback to stop training when certain loss was achieved
8
+ class StopAtLossCallback(TrainerCallback):
9
+ def __init__(self, threshold):
10
+ self.threshold = threshold
11
+
12
+ def on_log(self, args, state, control, logs=None, **kwargs):
13
+ if logs and "loss" in logs:
14
+ current_loss = logs["loss"]
15
+ if current_loss <= self.threshold:
16
+ print(f"\n[!] Target loss achieved. Training stopped.")
17
+ control.should_training_stop = True
18
+
19
+ # Model config
20
+ max_seq_length = 2048
21
+ model_name = "Qwen/Qwen2.5-VL-3B-Instruct"
22
+
23
+ model, processor = FastVisionModel.from_pretrained(
24
+ model_name = model_name,
25
+ max_seq_length = max_seq_length,
26
+ dtype = torch.bfloat16,
27
+ load_in_4bit = False, # Jeśli braknie VRAM, zmień na True!
28
+ )
29
+
30
+ tokenizer = processor.tokenizer
31
+
32
+ # Add new tokens
33
+ new_tokens = ["<move>", "</move>"]
34
+ for lvl in range(4):
35
+ for val in range(1024):
36
+ new_tokens.append(f"<m_{lvl}_{val}>")
37
+
38
+ tokenizer.add_special_tokens({'additional_special_tokens': new_tokens})
39
+ model.resize_token_embeddings(len(tokenizer))
40
+
41
+ # Train everything except vision
42
+ for name, param in model.named_parameters():
43
+ # Jeśli w nazwie warstwy jest 'visual' lub 'vision' - nie trenujemy jej
44
+ if "visual" in name or "vision" in name:
45
+ param.requires_grad = False
46
+ else:
47
+ param.requires_grad = True
48
+
49
+ model.gradient_checkpointing_enable()
50
+
51
+ # Dataset
52
+ dataset = load_dataset("json", data_files="la_dataset.jsonl", split="train")
53
+
54
+ def format_qwen_chat(examples):
55
+ texts = []
56
+ for instruction, output in zip(examples["instruction"], examples["output"]):
57
+ chat = [
58
+ {"role": "system", "content": "You are an embodied AI. You reason about your physical state and output precise motor actions inside <move></move> tags."},
59
+ {"role": "user", "content": instruction},
60
+ {"role": "assistant", "content": output}
61
+ ]
62
+ text = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=False)
63
+ texts.append(text)
64
+ return { "text" : texts }
65
+
66
+ dataset = dataset.map(format_qwen_chat, batched = True)
67
+
68
+ response_template = "<|im_start|>assistant\n"
69
+ collator = DataCollatorForCompletionOnlyLM(response_template=response_template, tokenizer=tokenizer)
70
+
71
+ # Training
72
+ trainer = SFTTrainer(
73
+ model = model,
74
+ tokenizer = tokenizer,
75
+ train_dataset = dataset,
76
+ dataset_text_field = "text",
77
+ max_seq_length = max_seq_length,
78
+ dataset_num_proc = 2,
79
+ data_collator = collator,
80
+ callbacks=[StopAtLossCallback(threshold=1.0)], # target loss
81
+ args = TrainingArguments(
82
+ per_device_train_batch_size = 2,
83
+ gradient_accumulation_steps = 8,
84
+ warmup_steps = 100,
85
+ num_train_epochs = 100,
86
+ learning_rate = 2e-5,
87
+ fp16 = not torch.cuda.is_bf16_supported(),
88
+ bf16 = torch.cuda.is_bf16_supported(),
89
+ logging_steps = 10,
90
+ output_dir = "outputs_qwen_vl_fft",
91
+ optim = "adamw_8bit",
92
+ save_strategy="no",
93
+ save_total_limit=1,
94
+ ),
95
+ )
96
+
97
+ trainer.train()
98
+
99
+ new_save_path = "Qwen2_5_VL_VLA_Final"
100
+ model.save_pretrained(new_save_path)
101
+ processor.save_pretrained(new_save_path)
102
+
103
+ print(f"Training done, model saved to {new_save_path}.")
chat_template.jinja ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system
2
+ You are a helpful assistant.<|im_end|>
3
+ {% endif %}<|im_start|>{{ message['role'] }}
4
+ {% if message['content'] is string %}{{ message['content'] }}<|im_end|>
5
+ {% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>
6
+ {% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
7
+ {% endif %}
config.json ADDED
@@ -0,0 +1,116 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "Qwen2_5_VLForConditionalGeneration"
4
+ ],
5
+ "bos_token_id": null,
6
+ "dtype": "bfloat16",
7
+ "eos_token_id": 151645,
8
+ "image_token_id": 151655,
9
+ "model_name": "unsloth/Qwen2.5-VL-3B-Instruct",
10
+ "model_type": "qwen2_5_vl",
11
+ "pad_token_id": 151654,
12
+ "text_config": {
13
+ "attention_dropout": 0.0,
14
+ "bos_token_id": 151643,
15
+ "dtype": "bfloat16",
16
+ "eos_token_id": 151645,
17
+ "hidden_act": "silu",
18
+ "hidden_size": 2048,
19
+ "initializer_range": 0.02,
20
+ "intermediate_size": 11008,
21
+ "layer_types": [
22
+ "full_attention",
23
+ "full_attention",
24
+ "full_attention",
25
+ "full_attention",
26
+ "full_attention",
27
+ "full_attention",
28
+ "full_attention",
29
+ "full_attention",
30
+ "full_attention",
31
+ "full_attention",
32
+ "full_attention",
33
+ "full_attention",
34
+ "full_attention",
35
+ "full_attention",
36
+ "full_attention",
37
+ "full_attention",
38
+ "full_attention",
39
+ "full_attention",
40
+ "full_attention",
41
+ "full_attention",
42
+ "full_attention",
43
+ "full_attention",
44
+ "full_attention",
45
+ "full_attention",
46
+ "full_attention",
47
+ "full_attention",
48
+ "full_attention",
49
+ "full_attention",
50
+ "full_attention",
51
+ "full_attention",
52
+ "full_attention",
53
+ "full_attention",
54
+ "full_attention",
55
+ "full_attention",
56
+ "full_attention",
57
+ "full_attention"
58
+ ],
59
+ "max_position_embeddings": 128000,
60
+ "max_window_layers": 70,
61
+ "model_type": "qwen2_5_vl_text",
62
+ "num_attention_heads": 16,
63
+ "num_hidden_layers": 36,
64
+ "num_key_value_heads": 2,
65
+ "pad_token_id": 151654,
66
+ "rms_norm_eps": 1e-06,
67
+ "rope_parameters": {
68
+ "mrope_section": [
69
+ 16,
70
+ 24,
71
+ 24
72
+ ],
73
+ "rope_theta": 1000000.0,
74
+ "rope_type": "default",
75
+ "type": "default"
76
+ },
77
+ "sliding_window": null,
78
+ "use_cache": true,
79
+ "use_sliding_window": false,
80
+ "vocab_size": 155763
81
+ },
82
+ "tie_word_embeddings": true,
83
+ "transformers_version": "5.5.0",
84
+ "unsloth_fixed": true,
85
+ "unsloth_version": "2026.6.1",
86
+ "use_cache": false,
87
+ "video_token_id": 151656,
88
+ "vision_config": {
89
+ "depth": 32,
90
+ "dtype": "bfloat16",
91
+ "fullatt_block_indexes": [
92
+ 7,
93
+ 15,
94
+ 23,
95
+ 31
96
+ ],
97
+ "hidden_act": "silu",
98
+ "hidden_size": 1280,
99
+ "in_channels": 3,
100
+ "in_chans": 3,
101
+ "initializer_range": 0.02,
102
+ "intermediate_size": 3420,
103
+ "model_type": "qwen2_5_vl",
104
+ "num_heads": 16,
105
+ "out_hidden_size": 2048,
106
+ "patch_size": 14,
107
+ "spatial_merge_size": 2,
108
+ "spatial_patch_size": 14,
109
+ "temporal_patch_size": 2,
110
+ "tokens_per_second": 2,
111
+ "window_size": 112
112
+ },
113
+ "vision_end_token_id": 151653,
114
+ "vision_start_token_id": 151652,
115
+ "vision_token_id": 151654
116
+ }
generated_motion.gif ADDED

Git LFS Details

  • SHA256: d7d3456f64fff6392791091d97307cade4175d99f0e68f80c1ee25f0095842d4
  • Pointer size: 131 Bytes
  • Size of remote file: 459 kB
generation_config.json ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "do_sample": true,
3
+ "eos_token_id": [
4
+ 151645,
5
+ 151645,
6
+ 151643
7
+ ],
8
+ "max_length": 128000,
9
+ "pad_token_id": 151654,
10
+ "repetition_penalty": 1.05,
11
+ "temperature": 1e-06,
12
+ "transformers_version": "5.5.0"
13
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b4b58a70470c4107606f640f34b4c3846ee936e0ee291d49fbf8f8adddae26af
3
+ size 7525015808
processor_config.json ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "image_processor": {
3
+ "do_convert_rgb": true,
4
+ "do_normalize": true,
5
+ "do_rescale": true,
6
+ "do_resize": true,
7
+ "image_mean": [
8
+ 0.48145466,
9
+ 0.4578275,
10
+ 0.40821073
11
+ ],
12
+ "image_processor_type": "Qwen2VLImageProcessor",
13
+ "image_std": [
14
+ 0.26862954,
15
+ 0.26130258,
16
+ 0.27577711
17
+ ],
18
+ "merge_size": 2,
19
+ "patch_size": 14,
20
+ "resample": 3,
21
+ "rescale_factor": 0.00392156862745098,
22
+ "size": {
23
+ "longest_edge": 12845056,
24
+ "shortest_edge": 3136
25
+ },
26
+ "temporal_patch_size": 2
27
+ },
28
+ "processor_class": "Qwen2_5_VLProcessor",
29
+ "video_processor": {
30
+ "do_convert_rgb": true,
31
+ "do_normalize": true,
32
+ "do_rescale": true,
33
+ "do_resize": true,
34
+ "do_sample_frames": false,
35
+ "image_mean": [
36
+ 0.48145466,
37
+ 0.4578275,
38
+ 0.40821073
39
+ ],
40
+ "image_std": [
41
+ 0.26862954,
42
+ 0.26130258,
43
+ 0.27577711
44
+ ],
45
+ "max_frames": 768,
46
+ "merge_size": 2,
47
+ "min_frames": 4,
48
+ "patch_size": 14,
49
+ "resample": 3,
50
+ "rescale_factor": 0.00392156862745098,
51
+ "return_metadata": false,
52
+ "size": {
53
+ "longest_edge": 12845056,
54
+ "shortest_edge": 3136
55
+ },
56
+ "temporal_patch_size": 2,
57
+ "video_processor_type": "Qwen2VLVideoProcessor"
58
+ }
59
+ }
rvq_model/Mean.npy ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:26e136555dab04c94a129d446c26e6b9939cbf045fbf77bcf5462c1fb5a2001c
3
+ size 1180
rvq_model/Std.npy ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6565a65ed9b31e23c328829a309e1c482be8b85fd23b43d65451a9b19a917f40
3
+ size 1180
rvq_model/__pycache__/rvq_model.cpython-311.pyc ADDED
Binary file (10.4 kB). View file
 
rvq_model/motion_rvq_weights.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7bbee9c0dfbe42875ace86db0f584f8634e83711f231bbaca551eb65a3e504f8
3
+ size 74587252
rvq_model/rvq_model.py ADDED
@@ -0,0 +1,140 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.nn.functional as F
4
+
5
+
6
+ class EMAVectorQuantizer(nn.Module):
7
+ def __init__(
8
+ self,
9
+ num_embeddings=512,
10
+ embedding_dim=256,
11
+ commitment_cost=0.25,
12
+ decay=0.99,
13
+ epsilon=1e-5,
14
+ ):
15
+ super().__init__()
16
+ self.num_embeddings = num_embeddings
17
+ self.embedding_dim = embedding_dim
18
+ self.commitment_cost = commitment_cost
19
+ self.decay = decay
20
+ self.epsilon = epsilon
21
+
22
+ embed = torch.randn(num_embeddings, embedding_dim)
23
+ self.register_buffer("embedding", embed)
24
+ self.register_buffer("cluster_size", torch.zeros(num_embeddings))
25
+ self.register_buffer("ema_w", embed.clone())
26
+
27
+ def forward(self, z):
28
+ z = z.permute(0, 2, 1).contiguous()
29
+ z_flattened = z.view(-1, self.embedding_dim)
30
+
31
+ distances = (
32
+ torch.sum(z_flattened**2, dim=1, keepdim=True)
33
+ + torch.sum(self.embedding**2, dim=1)
34
+ - 2 * torch.matmul(z_flattened, self.embedding.t())
35
+ )
36
+
37
+ min_encoding_indices = torch.argmin(distances, dim=1)
38
+ z_q = F.embedding(min_encoding_indices, self.embedding)
39
+
40
+ if self.training:
41
+ encodings = F.one_hot(min_encoding_indices, self.num_embeddings).float()
42
+ self.cluster_size.data.mul_(self.decay).add_(encodings.sum(0), alpha=1 - self.decay)
43
+
44
+ n = self.cluster_size.sum()
45
+ cluster_size = (self.cluster_size + self.epsilon) / (
46
+ n + self.num_embeddings * self.epsilon
47
+ ) * n
48
+
49
+ dw = torch.matmul(encodings.t(), z_flattened)
50
+ self.ema_w.data.mul_(self.decay).add_(dw, alpha=1 - self.decay)
51
+ self.embedding.data.copy_(self.ema_w / cluster_size.unsqueeze(1))
52
+
53
+ loss = self.commitment_cost * F.mse_loss(z_q.detach(), z_flattened)
54
+ z_q = z_flattened + (z_q - z_flattened).detach()
55
+ z_q = z_q.view(z.shape).permute(0, 2, 1).contiguous()
56
+
57
+ return z_q, min_encoding_indices.view(z.shape[0], z.shape[1]), loss
58
+
59
+
60
+ class RVQ(nn.Module):
61
+ def __init__(self, num_levels=3, num_embeddings=512, embedding_dim=256):
62
+ super().__init__()
63
+ self.num_levels = num_levels
64
+ self.quantizers = nn.ModuleList(
65
+ [EMAVectorQuantizer(num_embeddings, embedding_dim) for _ in range(num_levels)]
66
+ )
67
+
68
+ def forward(self, z):
69
+ quantized_out = 0
70
+ residual = z
71
+ all_indices = []
72
+ total_loss = 0
73
+
74
+ for quantizer in self.quantizers:
75
+ z_q, indices, loss = quantizer(residual)
76
+ quantized_out = quantized_out + z_q
77
+ residual = residual - z_q
78
+ all_indices.append(indices)
79
+ total_loss += loss
80
+
81
+ return quantized_out, torch.stack(all_indices, dim=1), total_loss
82
+
83
+
84
+ class ResBlock1D(nn.Module):
85
+ def __init__(self, channels):
86
+ super().__init__()
87
+ self.net = nn.Sequential(
88
+ nn.Conv1d(channels, channels, kernel_size=3, padding=1),
89
+ nn.LeakyReLU(0.2, inplace=True),
90
+ nn.Conv1d(channels, channels, kernel_size=3, padding=1),
91
+ )
92
+
93
+ def forward(self, x):
94
+ return x + self.net(x)
95
+
96
+
97
+ class MotionEncoder(nn.Module):
98
+ def __init__(self, in_channels=263, latent_dim=512):
99
+ super().__init__()
100
+ self.net = nn.Sequential(
101
+ nn.Conv1d(in_channels, 512, kernel_size=3, padding=1),
102
+ nn.LeakyReLU(0.2, inplace=True),
103
+ ResBlock1D(512),
104
+ ResBlock1D(512),
105
+ ResBlock1D(512),
106
+ nn.Conv1d(512, latent_dim, kernel_size=8, stride=4, padding=2),
107
+ )
108
+
109
+ def forward(self, x):
110
+ return self.net(x)
111
+
112
+
113
+ class MotionDecoder(nn.Module):
114
+ def __init__(self, latent_dim=512, out_channels=263):
115
+ super().__init__()
116
+ self.net = nn.Sequential(
117
+ nn.ConvTranspose1d(latent_dim, 512, kernel_size=8, stride=4, padding=2),
118
+ nn.LeakyReLU(0.2, inplace=True),
119
+ ResBlock1D(512),
120
+ ResBlock1D(512),
121
+ ResBlock1D(512),
122
+ nn.Conv1d(512, out_channels, kernel_size=3, padding=1),
123
+ )
124
+
125
+ def forward(self, z_q):
126
+ return self.net(z_q)
127
+
128
+
129
+ class MotionRVQ_VAE(nn.Module):
130
+ def __init__(self):
131
+ super().__init__()
132
+ self.encoder = MotionEncoder(in_channels=263, latent_dim=512)
133
+ self.rvq = RVQ(num_levels=4, num_embeddings=1024, embedding_dim=512)
134
+ self.decoder = MotionDecoder(latent_dim=512, out_channels=263)
135
+
136
+ def forward(self, x):
137
+ z = self.encoder(x)
138
+ z_q, token_indices, commitment_loss = self.rvq(z)
139
+ x_recon = self.decoder(z_q)
140
+ return x_recon, token_indices, commitment_loss
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:436b4575068de24aefa74f606e5c1b82a604c169a3502384d8cfd68761ed3a4a
3
+ size 12183771
tokenizer_config.json ADDED
The diff for this file is too large to render. See raw diff
 
walker.jpeg ADDED

Git LFS Details

  • SHA256: ef13a33a6ea0a1a41435e22a6778df2e8e6d3a59a37f1fcc8635424bd778e1f1
  • Pointer size: 131 Bytes
  • Size of remote file: 285 kB