初始化项目,由ModelHub XC社区提供模型
Model: EmpathicRobotics/vla-1.7b-qwen3-v2 Source: Original Platform
This commit is contained in:
104
tools/encode/encode_agent.py
Normal file
104
tools/encode/encode_agent.py
Normal file
@@ -0,0 +1,104 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Agent (3D pose) encoder -- turns a real 17-joint pose window into `<agent>`
|
||||
tokens, the reverse of tools/eval/decode_agent_tokens.py. Reuses
|
||||
pipeline_pose/phase5_adaptive_pchip.py's build_token_str() verbatim (pure
|
||||
numpy, no cluster dependency at all) rather than re-deriving the adaptive
|
||||
PCHIP control-point selection.
|
||||
|
||||
Use case: you have a REAL 3D pose sequence (motion capture, or your own
|
||||
video run through an HRNet+MotionBERT-style pipeline) and want the model to
|
||||
continue/predict from it -- this is exactly the "agent completion" behavior
|
||||
already verified for this model (give a partial <agent> block, it completes
|
||||
all 17 joints). Unlike seed2/cosmos/snac, this is the only encoder where
|
||||
"raw input" isn't a stock media file -- it's already-estimated 3D joint
|
||||
positions, which is a fair thing to require (you can't derive metric 3D pose
|
||||
from nothing; some upstream pose-estimation step is unavoidable no matter
|
||||
who's doing the encoding).
|
||||
|
||||
Input contract (IMPORTANT, easy to get wrong):
|
||||
- shape (8, 17, 3) float -- exactly 8 frames (this model's WINDOW_FRAMES),
|
||||
NOT 24 (that's the newer 2026-07-23 pipeline convention this model never
|
||||
saw), 17 joints in the exact order below, xyz in METRES.
|
||||
- ROOT-CENTERED: pelvis (joint 0) must be at [0,0,0] in every frame --
|
||||
subtract the pelvis position from all 17 joints per-frame yourself first
|
||||
if your source data isn't already root-relative (see
|
||||
pipeline_pose/phase3_kinematics_processor.py's split_root_motion() for
|
||||
the exact convention this project uses).
|
||||
- Values should stay within [-2.0, +2.0]m per axis (COORD_RANGE) --
|
||||
quantize() clips silently outside that range, so a badly-scaled pose
|
||||
(e.g. millimetres instead of metres) will silently flatten to the
|
||||
boundary rather than erroring. No automatic unit detection is attempted.
|
||||
|
||||
Joint order:
|
||||
pelvis, r_hip, r_knee, r_ankle, l_hip, l_knee, l_ankle, spine, thorax,
|
||||
nose, head_top, l_shoulder, l_elbow, l_wrist, r_shoulder, r_elbow, r_wrist
|
||||
|
||||
Usage:
|
||||
python tools/encode/encode_agent.py --input pose.json
|
||||
# pose.json: {"states": [[[x,y,z], ...17 joints...], ...8 frames...]}
|
||||
python tools/encode/encode_agent.py --input pose.npy
|
||||
# pose.npy: numpy array, shape (8, 17, 3)
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "pipeline_pose"))
|
||||
from phase5_adaptive_pchip import build_token_str, JOINT_NAMES, N_JOINTS, TARGET_FPS, COORD_RANGE # noqa: E402
|
||||
|
||||
WINDOW_FRAMES = 8 # this model's convention -- NOT the newer 24-frame pipeline
|
||||
|
||||
|
||||
def load_states(path: str):
|
||||
import numpy as np
|
||||
|
||||
if path.endswith(".npy"):
|
||||
states = np.load(path)
|
||||
else:
|
||||
with open(path) as f:
|
||||
data = json.load(f)
|
||||
states = np.array(data["states"], dtype=np.float32)
|
||||
|
||||
if states.shape != (WINDOW_FRAMES, N_JOINTS, 3):
|
||||
raise ValueError(
|
||||
f"Expected shape ({WINDOW_FRAMES}, {N_JOINTS}, 3), got {states.shape}. "
|
||||
f"This model was trained on 8-frame windows -- 24-frame input (the newer "
|
||||
f"pipeline convention) will NOT tokenize correctly here."
|
||||
)
|
||||
|
||||
pelvis = states[:, 0, :]
|
||||
if not (abs(pelvis).max() < 1e-4):
|
||||
print(f"WARNING: pelvis (joint 0) is not at origin (max |pelvis|={abs(pelvis).max():.4f}m) -- "
|
||||
f"auto-centering now. If this wasn't intended, check your source data's convention.",
|
||||
file=sys.stderr)
|
||||
states = states - pelvis[:, None, :]
|
||||
|
||||
bad = states[abs(states) > COORD_RANGE]
|
||||
if bad.size > 0:
|
||||
print(f"WARNING: {bad.size} coordinate value(s) outside [-{COORD_RANGE}, {COORD_RANGE}]m -- "
|
||||
f"will be silently clipped by quantize(). Check units (expected metres).",
|
||||
file=sys.stderr)
|
||||
|
||||
return states
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--input", required=True, help=".json (with a 'states' key) or .npy file, shape (8,17,3)")
|
||||
ap.add_argument("--fps", type=int, default=TARGET_FPS)
|
||||
args = ap.parse_args()
|
||||
|
||||
states = load_states(args.input)
|
||||
token_str, cp_counts = build_token_str(states, fps=args.fps)
|
||||
|
||||
print(f"Encoded {N_JOINTS} joints, {sum(cp_counts.values())} total control points:")
|
||||
for name in JOINT_NAMES:
|
||||
print(f" {name}: {cp_counts[name]} CPs")
|
||||
print()
|
||||
print("<agent> " + token_str + " </agent>")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
95
tools/encode/encode_cosmos.py
Normal file
95
tools/encode/encode_cosmos.py
Normal file
@@ -0,0 +1,95 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Cosmos video-chunk encoder -- turns 8 real video frames into `<cosmos_N>`
|
||||
tokens (200 raw ids, no offset), the reverse of tools/decode/decode_cosmos.py.
|
||||
Reproduces the OLD (vla-1.7b-qwen3-v2 training-era) preprocessing convention
|
||||
EXACTLY, not the current/newer aspect-preserving pipeline convention
|
||||
(2026-07-23 pivot, 896 tokens/chunk) -- recovered from git history
|
||||
(commit edf25393, before 38d8e5f2 switched to aspect-preserving):
|
||||
|
||||
Resize((160,160)) # direct squash/stretch to 160x160, NOT an
|
||||
# aspect-preserving crop -- distorts aspect
|
||||
# ratio on purpose (this model's own training
|
||||
# convention, don't "fix" it here)
|
||||
ToTensor()
|
||||
Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]) # -> range [-1, 1]
|
||||
|
||||
8 frames -> stack -> permute(1,0,2,3) -> (3,8,160,160) -> unsqueeze(0) ->
|
||||
(1,3,8,160,160) -> CausalVideoTokenizer.encode() -> (1,2,10,10) == 200 ids,
|
||||
checkpoint nvidia/Cosmos-Tokenizer-DV8x16x16 (encoder.jit -- same repo
|
||||
decode_cosmos.py already downloads decoder.jit from).
|
||||
|
||||
Usage:
|
||||
python tools/encode/encode_cosmos.py --frames f0.png f1.png ... f7.png --output tokens.txt
|
||||
# exactly 8 frame image paths, in temporal order
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "decode", "vendor"))
|
||||
|
||||
TARGET_SIZE = 160
|
||||
N_FRAMES = 8
|
||||
COSMOS_HF_REPO = "nvidia/Cosmos-Tokenizer-DV8x16x16"
|
||||
_LOCAL_CHECKPOINT_ENC = "/e/project1/reformo/nguyen38/prototype/pretrained_ckpts/Cosmos-Tokenizer-DV8x16x16/encoder.jit"
|
||||
|
||||
|
||||
def _resolve_checkpoint_enc() -> str:
|
||||
if os.path.exists(_LOCAL_CHECKPOINT_ENC):
|
||||
return _LOCAL_CHECKPOINT_ENC
|
||||
from huggingface_hub import hf_hub_download
|
||||
print(f"Local checkpoint not found -- downloading encoder.jit from {COSMOS_HF_REPO} "
|
||||
f"(~350MB, cached for future runs)...")
|
||||
return hf_hub_download(repo_id=COSMOS_HF_REPO, filename="encoder.jit")
|
||||
|
||||
|
||||
def encode_frames(frame_paths: list) -> list:
|
||||
if len(frame_paths) != N_FRAMES:
|
||||
raise ValueError(f"Expected exactly {N_FRAMES} frame paths, got {len(frame_paths)}")
|
||||
|
||||
import torch
|
||||
import torchvision.transforms as T
|
||||
from PIL import Image
|
||||
from cosmos_tokenizer.video_lib import CausalVideoTokenizer
|
||||
|
||||
transform = T.Compose([
|
||||
T.Resize((TARGET_SIZE, TARGET_SIZE)),
|
||||
T.ToTensor(),
|
||||
T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
])
|
||||
frames = [transform(Image.open(p).convert("RGB")) for p in frame_paths]
|
||||
video = torch.stack(frames, dim=0) # (T, 3, H, W)
|
||||
video = video.permute(1, 0, 2, 3) # (3, T, H, W)
|
||||
video = video.unsqueeze(0) # (1, 3, T, H, W)
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
enc = CausalVideoTokenizer(checkpoint_enc=_resolve_checkpoint_enc()).to(device)
|
||||
with torch.no_grad():
|
||||
indices = enc.encode(video.to(device))[0] # (1, 2, 10, 10)
|
||||
|
||||
ids = indices.reshape(-1).tolist()
|
||||
if len(ids) != 200:
|
||||
raise ValueError(f"Expected 200 raw ids, got {len(ids)} -- checkpoint/shape mismatch")
|
||||
return ids
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--frames", nargs=8, required=True, metavar="FRAME",
|
||||
help="Exactly 8 frame image paths, in temporal order")
|
||||
ap.add_argument("--output", help="Optional: write comma-separated ids to this file")
|
||||
args = ap.parse_args()
|
||||
|
||||
ids = encode_frames(args.frames)
|
||||
out = ",".join(str(i) for i in ids)
|
||||
print(f"{len(ids)} cosmos tokens:")
|
||||
print(out)
|
||||
if args.output:
|
||||
with open(args.output, "w") as f:
|
||||
f.write(out)
|
||||
print(f"Saved: {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
66
tools/encode/encode_seed2.py
Normal file
66
tools/encode/encode_seed2.py
Normal file
@@ -0,0 +1,66 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Seed2 encoder -- turns a real image into `<seed2_N>` tokens (0-8191, no
|
||||
offset), the reverse of tools/decode/decode_seed2.py. Reuses that module's
|
||||
_load_seed2_tokenizer() (same runtime patches: the transformers import move
|
||||
+ BertLMHeadModel.cls=None guard) rather than re-deriving them -- the public
|
||||
ontocord/seed2 repo's own seed2_tokenizer.py still has both bugs unpatched
|
||||
(verified 2026-07-23), so any fresh download needs these regardless of
|
||||
whether you're encoding or decoding.
|
||||
|
||||
Preprocessing: Seed2Tokenizer.encode_image() does its own internal resize to
|
||||
224x224 (CLIP-style Resize+Normalize, see seed2_tokenizer.py's `self.processor`)
|
||||
-- pass a PIL image straight through, no manual resize needed first.
|
||||
|
||||
Usage:
|
||||
python tools/encode/encode_seed2.py --image photo.jpg
|
||||
# prints 32 raw ids; wrap as <seed2> <seed2_N> ... </seed2> to splice
|
||||
# into a prompt for this model (v2's convention -- no offset needed)
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "decode"))
|
||||
from decode_seed2 import _load_seed2_tokenizer, NUM_IMAGE_TOKENS # noqa: E402
|
||||
|
||||
|
||||
def encode_image(image_path: str) -> list:
|
||||
from PIL import Image
|
||||
|
||||
# _load_seed2_tokenizer() os.chdir()s -- resolve a relative image_path
|
||||
# against the original cwd *before* that happens, same class of bug
|
||||
# already fixed once in decode_seed2.py's own --output handling.
|
||||
image_path = os.path.abspath(image_path)
|
||||
|
||||
Seed2Tokenizer, seed2_dir = _load_seed2_tokenizer()
|
||||
tokenizer = Seed2Tokenizer.from_pretrained(seed2_dir).eval()
|
||||
if hasattr(tokenizer, "cuda") and __import__("torch").cuda.is_available():
|
||||
tokenizer = tokenizer.cuda()
|
||||
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
ids = tokenizer.encode_image(image_pil=image)
|
||||
ids = ids.view(-1).tolist()
|
||||
|
||||
bad = [t for t in ids if not (0 <= t < NUM_IMAGE_TOKENS)]
|
||||
if bad:
|
||||
raise ValueError(f"encode_image produced out-of-range ids: {bad[:5]}... (expected [0, {NUM_IMAGE_TOKENS}))")
|
||||
return ids
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--image", required=True, help="Path to a real image file")
|
||||
args = ap.parse_args()
|
||||
|
||||
ids = encode_image(args.image)
|
||||
print(f"{len(ids)} seed2 tokens:")
|
||||
print(",".join(str(i) for i in ids))
|
||||
print()
|
||||
print("As a prompt fragment:")
|
||||
print("<seed2> " + " ".join(f"<seed2_{i}>" for i in ids) + " </seed2>")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
78
tools/encode/encode_snac.py
Normal file
78
tools/encode/encode_snac.py
Normal file
@@ -0,0 +1,78 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
SNAC audio encoder -- turns a real audio/video file into `<snac_N>` tokens
|
||||
(listen format, 3 tokens/base-frame), the reverse of tools/decode/decode_snac.py.
|
||||
Reuses the exact encode_listen() logic already in pipeline_pose/snac_finevideo.py
|
||||
(unchanged since inception) rather than re-deriving it.
|
||||
|
||||
This model (vla-1.7b-qwen3-v2) only ever saw listen-format audio wrapped in
|
||||
the generic <snac> tag -- NOT the newer (2026-07-23) <listen>/<speak>
|
||||
convention or the speak-format L2 tokens. Output here matches that: always
|
||||
listen-format, always <snac> wrapper.
|
||||
|
||||
Usage:
|
||||
python tools/encode/encode_snac.py --input clip.wav --output tokens.txt
|
||||
python tools/encode/encode_snac.py --input video.mp4 --output tokens.txt
|
||||
# any format ffmpeg can read (audio extracted automatically, works on
|
||||
# video files too -- just uses the audio track)
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "pipeline_pose"))
|
||||
|
||||
SAMPLE_RATE = 24000
|
||||
SNAC_MODEL = "hubertsiuzdak/snac_24khz"
|
||||
|
||||
|
||||
def _get_ffmpeg() -> str:
|
||||
import imageio_ffmpeg
|
||||
return imageio_ffmpeg.get_ffmpeg_exe()
|
||||
|
||||
|
||||
def extract_audio(input_path: str):
|
||||
import numpy as np
|
||||
|
||||
cmd = [_get_ffmpeg(), "-y", "-i", input_path, "-vn", "-ac", "1", "-ar", str(SAMPLE_RATE), "-f", "f32le", "-"]
|
||||
result = subprocess.run(cmd, capture_output=True, timeout=300)
|
||||
if result.returncode != 0 or not result.stdout:
|
||||
raise RuntimeError(f"ffmpeg failed to extract audio from {input_path}: {result.stderr.decode(errors='replace')[:500]}")
|
||||
audio = np.frombuffer(result.stdout, dtype=np.float32).copy()
|
||||
if len(audio) == 0:
|
||||
raise RuntimeError(f"No audio extracted from {input_path} -- does it have an audio track?")
|
||||
return audio
|
||||
|
||||
|
||||
def encode_file(input_path: str) -> list:
|
||||
import torch
|
||||
from snac import SNAC
|
||||
from snac_finevideo import encode_listen # reused verbatim, not re-derived
|
||||
|
||||
audio = extract_audio(input_path)
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
model = SNAC.from_pretrained(SNAC_MODEL).eval().to(device)
|
||||
tokens = encode_listen(audio, model, device) # already "<snac_N>" strings, listen-format
|
||||
return tokens
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--input", required=True, help="Audio or video file (any ffmpeg-readable format)")
|
||||
ap.add_argument("--output", help="Optional: write the <snac> block to this file")
|
||||
args = ap.parse_args()
|
||||
|
||||
tokens = encode_file(args.input)
|
||||
duration_s = len(tokens) / 3 / 12.5
|
||||
print(f"{len(tokens)} snac tokens ({len(tokens) // 3} base frames, ~{duration_s:.2f}s @ 12.5Hz base rate)")
|
||||
block = "<snac> " + " ".join(tokens) + " </snac>"
|
||||
print(block)
|
||||
if args.output:
|
||||
with open(args.output, "w") as f:
|
||||
f.write(block)
|
||||
print(f"Saved: {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user