mirror of
https://github.com/storytold/effectcraft.git
synced 2026-10-31 15:48:33 +00:00
Kristian MyramandClaude Sonnet 5.5
e4f4fea0f2
Mask down-sampler, feature fuser and projection, sine position encoding, occlusion embedding and bfloat16 rounding of stored memories. Max abs error 3.5e-6 on the memory features, 6e-8 on the position encoding. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
113 lines
4.9 KiB
Python
113 lines
4.9 KiB
Python
#!/usr/bin/env python
|
|
"""Dump SAM 2 video reference tensors (CPU, float32) for the Rust port's reference tests.
|
|
|
|
Usage: python -I tools/sam2_video_reference.py MODEL_DIR OUT.safetensors
|
|
Runs `Sam2VideoModel` on a 7-frame synthetic clip (a textured disc moving right, hidden behind a bar on
|
|
frames 4 and 5), prompted with one click on frame 2, then tracks forward (2..6) and backward (2..0).
|
|
Saves the preprocessed frames, every frame's low-resolution mask and object score, the stored memory
|
|
features and object pointers, and the inputs/outputs of the memory encoder and memory attention.
|
|
"""
|
|
import sys
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from safetensors.torch import save_file
|
|
from transformers import Sam2VideoModel
|
|
from transformers.models.sam2_video.modeling_sam2_video import Sam2VideoInferenceSession
|
|
|
|
T, SIZE, PROMPT_FRAME = 7, 256, 2
|
|
|
|
|
|
def clip():
|
|
rng = np.random.default_rng(3)
|
|
yy, xx = np.mgrid[0:SIZE, 0:SIZE].astype(np.float32)
|
|
bg = 0.3 + 0.2 * np.sin(xx / 9.0) * np.cos(yy / 7.0)
|
|
frames = []
|
|
for t in range(T):
|
|
cx, cy = 80 + 12 * t, 128
|
|
disc = np.hypot(xx - cx, yy - cy) < 34
|
|
img = np.stack([bg, bg * 0.8 + 0.1, bg * 0.6 + 0.3], -1)
|
|
tex = 0.85 + 0.1 * np.sin((xx - cx) / 5)
|
|
img[disc] = np.stack([tex, 0.4 + 0.1 * np.cos((yy - cy) / 5), 0.15 + 0 * xx], -1)[disc]
|
|
if t in (4, 5): # a bar hides the disc
|
|
img[:, int(cx) - 50 : int(cx) + 50] = 0.2
|
|
img += rng.normal(0, 0.01, img.shape).astype(np.float32)
|
|
frames.append((img.clip(0, 1) * 255).astype(np.uint8))
|
|
return frames, (80 + 12 * PROMPT_FRAME, 128)
|
|
|
|
|
|
def main(model_dir, out):
|
|
torch.manual_seed(0)
|
|
torch.set_num_threads(4)
|
|
model = Sam2VideoModel.from_pretrained(model_dir).eval()
|
|
frames, click = clip()
|
|
mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
|
|
std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
|
|
x = torch.from_numpy(np.stack(frames)).permute(0, 3, 1, 2).float() / 255.0
|
|
x = F.interpolate(x, size=(1024, 1024), mode="bilinear", align_corners=False)
|
|
pixels = (x - mean) / std
|
|
|
|
taps = {"clip.pixel_values": pixels}
|
|
first = {"enc": True, "attn": True}
|
|
|
|
def enc_hook(_m, args, kwargs, out):
|
|
if first["enc"]:
|
|
first["enc"] = False
|
|
vf = args[0] if args else kwargs["vision_features"]
|
|
masks = args[1] if len(args) > 1 else kwargs["masks"]
|
|
taps["mem.enc.vision_features"] = vf.detach().float().clone()
|
|
taps["mem.enc.masks"] = masks.detach().float().clone()
|
|
taps["mem.enc.out_features"] = out[0].detach().float().clone()
|
|
taps["mem.enc.out_pos"] = out[1].detach().float().clone()
|
|
|
|
def attn_hook(_m, args, kwargs, out):
|
|
if first["attn"]:
|
|
first["attn"] = False
|
|
for k in ("current_vision_features", "memory", "current_vision_position_embeddings", "memory_posision_embeddings"):
|
|
taps[f"mem.attn.{k}"] = kwargs[k].detach().float().clone()
|
|
taps["mem.attn.num_pointers"] = torch.tensor([float(kwargs["num_object_pointer_tokens"])])
|
|
taps["mem.attn.out"] = out.detach().float().clone()
|
|
|
|
model.memory_encoder.register_forward_hook(enc_hook, with_kwargs=True)
|
|
model.memory_attention.register_forward_hook(attn_hook, with_kwargs=True)
|
|
|
|
sess = Sam2VideoInferenceSession(video=pixels, video_height=SIZE, video_width=SIZE, dtype=torch.float32)
|
|
obj_idx = sess.obj_id_to_idx(1)
|
|
scale = 1024.0 / SIZE
|
|
sess.add_point_inputs(
|
|
obj_idx,
|
|
PROMPT_FRAME,
|
|
{
|
|
"point_coords": torch.tensor([[[[click[0] * scale, click[1] * scale]]]], dtype=torch.float32),
|
|
"point_labels": torch.tensor([[[1]]], dtype=torch.int32),
|
|
},
|
|
)
|
|
sess.obj_with_new_inputs = [1]
|
|
|
|
def record(tag, o):
|
|
taps[f"video.{tag}.{o.frame_idx}.pred_masks"] = o.pred_masks.detach().float().clone()
|
|
taps[f"video.{tag}.{o.frame_idx}.score"] = o.object_score_logits.detach().float().clone()
|
|
|
|
with torch.no_grad():
|
|
record("prompt", model(sess, frame_idx=PROMPT_FRAME))
|
|
for o in model.propagate_in_video_iterator(sess, start_frame_idx=PROMPT_FRAME):
|
|
record("fwd", o)
|
|
for o in model.propagate_in_video_iterator(sess, start_frame_idx=PROMPT_FRAME, reverse=True):
|
|
record("rev", o)
|
|
|
|
d = sess.output_dict_per_obj[obj_idx]
|
|
for kind in ("cond_frame_outputs", "non_cond_frame_outputs"):
|
|
for f, o in d[kind].items():
|
|
taps[f"state.{f}.maskmem_features"] = o["maskmem_features"].detach().float().clone()
|
|
taps[f"state.{f}.maskmem_pos_enc"] = o["maskmem_pos_enc"].detach().float().clone()
|
|
taps[f"state.{f}.object_pointer"] = o["object_pointer"].detach().float().clone()
|
|
save_file({k: v.contiguous() for k, v in taps.items()}, out)
|
|
print(f"wrote {len(taps)} tensors to {out}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) != 3:
|
|
sys.exit(__doc__)
|
|
main(sys.argv[1], sys.argv[2])
|