Upload Trace-Anything (transformers port, original code)

653a5eb verified 14 days ago

16.5 kB

# ruff: noqa
# fmt: off
"""Trace Anything, vendored verbatim from ByteDance-Seed/TraceAnything @ 54677b5e7bf11510c2e8c917a509988ad379f8eb (trace_anything/): the files below
concatenated, package-internal imports replaced by relative imports of the other vendored files (generated by scripts/model_ports/common/vendor_flat.py;
do not edit by hand). Code license: Apache-2.0."""
from .ta_heads import PixelHead, ScalarHead
from .ta_layers import Block, PositionGetter, RoPE2D, get_1d_sincos_pos_embed_from_grid, get_patch_embed


# ==============================================================================================================
# trace_anything/trace_anything.py
# ==============================================================================================================
# Copyright (c) 2025 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.

# trace_anything/trace_anything.py
"""
Minimal inference-only model for this repo.
Assumptions (pruned by asserts):
- encoder_type == 'croco'
- decoder_type == 'transformer'
- head_type == 'dpt'
- targeting_mechanism == 'bspline_conf'
- optional: whether_local (bool)
"""

import math
import time
from copy import deepcopy
from functools import partial
from typing import Dict, List

import torch
import torch.nn as nn
from einops import rearrange
import numpy as np


from contextlib import contextmanager
import time



# ======== B-spline  ========
PRECOMPUTED_KNOTS = {
    4:  torch.tensor([0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0], dtype=torch.float32),
    7:  torch.tensor([0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.5, 1.0, 1.0, 1.0, 1.0], dtype=torch.float32),
    10: torch.tensor([0.0, 0.0, 0.0, 0.0, 1/3, 1/3, 1/3, 2/3, 2/3, 2/3, 1.0, 1.0, 1.0, 1.0], dtype=torch.float32),
}

def _precompute_knot_differences(n_ctrl_pts, degree, knots):
    denom1 = torch.zeros(n_ctrl_pts, degree + 1, device=knots.device)
    denom2 = torch.zeros(n_ctrl_pts, degree + 1, device=knots.device)
    for k in range(degree + 1):
        for i in range(n_ctrl_pts):
            denom1[i, k] = knots[i + k] - knots[i] if i + k < len(knots) else 0.0
            denom2[i, k] = knots[i + k + 1] - knots[i + 1] if i + k + 1 < len(knots) else 1.0
    return denom1, denom2

PRECOMPUTED_DENOMS = {n: _precompute_knot_differences(n, 3, PRECOMPUTED_KNOTS[n]) for n in [4, 7, 10]}

def _compute_bspline_basis(n_ctrl_pts, degree, t_values, knots, denom1, denom2):
    N = t_values.size(0)
    basis = torch.zeros(N, n_ctrl_pts, degree + 1, device=t_values.device)
    t = t_values
    basis_k0 = torch.zeros(N, n_ctrl_pts, device=t.device)
    for i in range(n_ctrl_pts):
        if i == n_ctrl_pts - 1:
            basis_k0[:, i] = ((knots[i] <= t) & (t <= knots[i + 1])).float()
        else:
            basis_k0[:, i] = ((knots[i] <= t) & (t < knots[i + 1])).float()
    basis[:, :, 0] = basis_k0
    for k in range(1, degree + 1):
        basis_k = torch.zeros(N, n_ctrl_pts, device=t.device)
        for i in range(n_ctrl_pts):
            term1 = ((t - knots[i]) / denom1[i, k]) * basis[:, i, k-1] if denom1[i, k] > 0 else 0.0
            term2 = ((knots[i + k + 1] - t) / denom2[i, k]) * basis[:, i + 1, k-1] if (denom2[i, k] > 0 and i + 1 < n_ctrl_pts) else 0.0
            basis_k[:, i] = term1 + term2
        basis[:, :, k] = basis_k
    return basis[:, :, degree]

def evaluate_bspline_conf(ctrl_pts3d, ctrl_conf, t_values):
    """ctrl_pts3d:[N_ctrl,H,W,3], ctrl_conf:[N_ctrl,H,W], t_values:[T] -> (T,H,W,3),(T,H,W)"""
    n_ctrl_pts, H, W, _ = ctrl_pts3d.shape
    assert n_ctrl_pts in (4, 7, 10), f"unsupported n_ctrl_pts={n_ctrl_pts}"
    degree = 3
    knot_vector = PRECOMPUTED_KNOTS[n_ctrl_pts].to(ctrl_pts3d.device)
    denom1, denom2 = [d.to(ctrl_pts3d.device) for d in PRECOMPUTED_DENOMS[n_ctrl_pts]]
    ctrl_pts3d = ctrl_pts3d.permute(0, 3, 1, 2)         # [N,3,H,W]
    ctrl_conf  = ctrl_conf.unsqueeze(-1).permute(0, 3, 1, 2)  # [N,1,H,W]
    basis = _compute_bspline_basis(n_ctrl_pts, degree, t_values, knot_vector, denom1, denom2)  # [T,N]
    basis = basis.view(-1, n_ctrl_pts, 1, 1, 1)               # [T,N,1,1,1]
    pts3d_t = torch.sum(basis * ctrl_pts3d.unsqueeze(0), dim=1).permute(0, 2, 3, 1)  # [T,H,W,3]
    conf_t  = torch.sum(basis * ctrl_conf.unsqueeze(0),  dim=1).squeeze(1)           # [T,H,W]
    return pts3d_t, conf_t

# ======== Encoders（仅 CroCo） ========
class CroCoEncoder(nn.Module):
    def __init__(
        self,
        img_size=512, patch_size=16, patch_embed_cls="ManyAR_PatchEmbed",
        embed_dim=768, num_heads=12, depth=12, mlp_ratio=4,
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
        pos_embed="RoPE100", attn_implementation="pytorch_naive",
    ):
        super().__init__()
        assert pos_embed.startswith("RoPE"), f"pos_embed must start with RoPE*, got {pos_embed}"
        self.patch_embed = get_patch_embed(patch_embed_cls, img_size, patch_size, embed_dim)
        freq = float(pos_embed[len("RoPE"):])
        self.rope = RoPE2D(freq=freq)
        self.enc_blocks = nn.ModuleList([
            Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio,
                  qkv_bias=True, norm_layer=nn.LayerNorm, rope=self.rope,
                  attn_implementation=attn_implementation)
            for _ in range(depth)
        ])
        self.enc_norm = norm_layer(embed_dim)

    def forward(self, image, true_shape):
        x, pos = self.patch_embed(image, true_shape=true_shape)
        for blk in self.enc_blocks:
            x = blk(x, pos)
        x = self.enc_norm(x)
        return x, pos


# ======== Decoder ========
class TraceDecoder(nn.Module):
    def __init__(
        self,
        random_image_idx_embedding: bool,
        enc_embed_dim: int,
        embed_dim: int = 768,
        num_heads: int = 12,
        depth: int = 12,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = True,
        drop: float = 0.0,
        attn_drop: float = 0.0,
        attn_implementation: str = "pytorch_naive",
        attn_bias_for_inference_enabled=True,
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
    ):
        super().__init__()
        self.decoder_embed = nn.Linear(enc_embed_dim, embed_dim, bias=True)
        self.dec_blocks = nn.ModuleList([
            Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio,
                  qkv_bias=qkv_bias, drop=drop, attn_drop=attn_drop,
                  norm_layer=nn.LayerNorm, attn_implementation=attn_implementation,
                  attn_bias_for_inference_enabled=attn_bias_for_inference_enabled)
            for _ in range(depth)
        ])
        self.random_image_idx_embedding = random_image_idx_embedding
        self.register_buffer(
            "image_idx_emb",
            torch.from_numpy(get_1d_sincos_pos_embed_from_grid(embed_dim, np.arange(1000))).float(),
            persistent=False,
        )
        self.dec_norm = norm_layer(embed_dim)

    def _get_random_image_pos(self, encoded_feats, batch_size, num_views, max_image_idx, device):
        image_ids = torch.zeros(batch_size, num_views, dtype=torch.long)
        image_ids[:, 0] = 0
        per_forward_pass_seed = torch.randint(0, 2 ** 32, (1,)).item()
        per_rank_generator = torch.Generator().manual_seed(per_forward_pass_seed)
        for b in range(batch_size):
            random_ids = torch.randperm(max_image_idx, generator=per_rank_generator)[:num_views - 1] + 1
            image_ids[b, 1:] = random_ids
        image_ids = image_ids.to(device)
        image_pos_list = []
        for i in range(num_views):
            num_patches = encoded_feats[i].shape[1]
            pos_for_view = self.image_idx_emb[image_ids[:, i]].unsqueeze(1).repeat(1, num_patches, 1)
            image_pos_list.append(pos_for_view)
        return torch.cat(image_pos_list, dim=1)

    def forward(self, encoded_feats: List[torch.Tensor], positions: List[torch.Tensor],
                image_ids: torch.Tensor, image_timesteps: torch.Tensor):
        x = torch.cat(encoded_feats, dim=1)
        pos = torch.cat(positions, dim=1)
        outputs = [x]
        x = self.decoder_embed(x)

        if self.random_image_idx_embedding:
            image_pos = self._get_random_image_pos(
                encoded_feats=encoded_feats,
                batch_size=encoded_feats[0].shape[0],
                num_views=len(encoded_feats),
                max_image_idx=self.image_idx_emb.shape[0] - 1,
                device=x.device,
            )
        else:
            num_embeddings = self.image_idx_emb.shape[0]
            indices = (image_timesteps * (num_embeddings - 1)).long().view(-1)
            image_pos = torch.index_select(self.image_idx_emb, dim=0, index=indices)
            image_pos = image_pos.view(1, image_timesteps.shape[1], self.image_idx_emb.shape[1])

        x += image_pos
        for blk in self.dec_blocks:
            x = blk(x, pos)
            outputs.append(x)

        x = self.dec_norm(x)
        outputs[-1] = x
        return outputs, image_pos


# ======== main model ========
class TraceAnything(nn.Module):
    def __init__(self, *, encoder_args: Dict, decoder_args: Dict, head_args: Dict,
                 targeting_mechanism: str = "bspline_conf",
                 poly_degree: int = 10, whether_local: bool = False):
        super().__init__()

        assert targeting_mechanism == "bspline_conf", f"Only bspline_conf is supported now, got {targeting_mechanism}"
        assert encoder_args.get("encoder_type") == "croco"
        assert decoder_args.get("decoder_type", "transformer") in ("transformer", "fast3r")
        assert head_args.get("head_type") == "dpt"

        self.targeting_mechanism = targeting_mechanism
        self.poly_degree = int(poly_degree)
        self.whether_local = bool(whether_local or head_args.get("with_local_head", False))

        # build encoder / decoder
        enc_args = deepcopy(encoder_args); enc_args.pop("encoder_type", None)
        self.encoder = CroCoEncoder(**enc_args)

        dec_args = deepcopy(decoder_args); dec_args.pop("decoder_type", None)
        self.decoder = TraceDecoder(**dec_args)

        # build heads
        feature_dim = 256
        last_dim = feature_dim // 2
        ed = encoder_args["embed_dim"]; dd = decoder_args["embed_dim"]
        hooks = [0, decoder_args["depth"] * 2 // 4, decoder_args["depth"] * 3 // 4, decoder_args["depth"]]
        self.ds_head_time = ScalarHead(
            num_channels=1, feature_dim=feature_dim, last_dim=last_dim,
            hooks_idx=hooks, dim_tokens=[ed, dd, dd, dd], head_type="regression",
            patch_size=head_args["patch_size"],
        )
        out_nchan = (3 + bool(head_args["conf_mode"])) * self.poly_degree
        self.ds_head_track = PixelHead(
            num_channels=out_nchan, feature_dim=feature_dim, last_dim=last_dim,
            hooks_idx=hooks, dim_tokens=[ed, dd, dd, dd], head_type="regression",
            patch_size=head_args["patch_size"],
            depth_mode=head_args["depth_mode"], conf_mode=head_args["conf_mode"],
        )
        if self.whether_local:
            self.ds_head_local = PixelHead(
                num_channels=out_nchan, feature_dim=feature_dim, last_dim=last_dim,
                hooks_idx=hooks, dim_tokens=[ed, dd, dd, dd], head_type="regression",
                patch_size=head_args["patch_size"],
                depth_mode=head_args["depth_mode"], conf_mode=head_args["conf_mode"],
            )

        self.time_head  = self.ds_head_time
        self.track_head = self.ds_head_track
        if self.whether_local:
            self.local_head = self.ds_head_local

        self.max_parallel_views_for_head = 25

    def _encode_images(self, views, chunk_size=400):
        B = views[0]["img"].shape[0]
        same_shape = all(v["img"].shape == views[0]["img"].shape for v in views)
        if same_shape:
            imgs = torch.cat([v["img"] for v in views], dim=0)
            true_shapes = torch.cat([v.get("true_shape", torch.tensor(v["img"].shape[-2:])[None].repeat(B, 1)) for v in views], dim=0)
            feats_chunks, pos_chunks = [], []
            for s in range(0, imgs.shape[0], chunk_size):
                e = min(s + chunk_size, imgs.shape[0])
                f, p = self.encoder(imgs[s:e], true_shapes[s:e])
                feats_chunks.append(f); pos_chunks.append(p)
            feats = torch.cat(feats_chunks, dim=0); pos = torch.cat(pos_chunks, dim=0)
            encoded_feats = torch.split(feats, B, dim=0)
            positions    = torch.split(pos,   B, dim=0)
            shapes       = torch.split(true_shapes, B, dim=0)
        else:
            encoded_feats, positions, shapes = [], [], []
            for v in views:
                img = v["img"]
                true_shape = v.get("true_shape", torch.tensor(img.shape[-2:])[None].repeat(B, 1))
                f, p = self.encoder(img, true_shape)
                encoded_feats.append(f); positions.append(p); shapes.append(true_shape)
        return encoded_feats, positions, shapes

    @torch.no_grad()
    def forward(self, views, profiling: bool = False):
        # 1) encode
        encoded_feats, positions, shapes = self._encode_images(views)

        # 2) build time embedding
        num_images = len(views)
        B, P, D = encoded_feats[0].shape
        image_ids, image_times = [], []
        for i, ef in enumerate(encoded_feats):
            num_patches = ef.shape[1]
            image_ids.extend([i] * num_patches)
            image_times.extend([views[i]["time_step"]] * num_patches)
        image_ids  = torch.tensor(image_ids * B,  device=encoded_feats[0].device).reshape(B, -1)
        image_times = torch.tensor(image_times * B, device=encoded_feats[0].device).reshape(B, -1)

        # 3) decode
        dec_output, _ = self.decoder(encoded_feats, positions, image_ids, image_times)

        # 4) gather outputs per view
        P_patches = P
        gathered_outputs_list = []
        for layer_output in dec_output:
            layer_output = rearrange(layer_output, 'B (n P) D -> (n B) P D', n=num_images, P=P_patches)
            gathered_outputs_list.append(layer_output)

        # 5) heads
        time_step = self.time_head(gathered_outputs_list, torch.stack(shapes))          # [N, 1] per view
        track_tmp = self.track_head(gathered_outputs_list, torch.stack(shapes))        # dict with pts3d/conf after postprocess
        ctrl_pts3d, ctrl_conf = track_tmp['pts3d'], track_tmp['conf']                  # [K,N,H,W,3], [K,N,H,W]

        if self.whether_local:
            local_tmp = self.local_head(gathered_outputs_list, torch.stack(shapes))
            ctrl_pts3d_local, ctrl_conf_local = local_tmp['pts3d'], local_tmp['conf']

        # 6) evaluate bspline track over all reference times
        results = [{} for _ in range(num_images)]
        t_values = torch.stack([time_step[i].squeeze() for i in range(num_images)])    # [N]
        for img_id in range(num_images):
            pts3d_t, conf_t = evaluate_bspline_conf(ctrl_pts3d[:, img_id], ctrl_conf[:, img_id], t_values.detach())
            res = {
                'time': time_step[img_id],
                'ctrl_pts3d': ctrl_pts3d[:, img_id],
                'ctrl_conf':  ctrl_conf[:,  img_id],
                'track_pts3d': [pts3d_t[[t]] for t in range(num_images)],
                'track_conf':  [conf_t[[t]]  for t in range(num_images)],
            }
            if self.whether_local:
                pts3d_t_l, conf_t_l = evaluate_bspline_conf(ctrl_pts3d_local[:, img_id], ctrl_conf_local[:, img_id], t_values.detach())
                res.update({
                    'ctrl_pts3d_local': ctrl_pts3d_local[:, img_id],
                    'ctrl_conf_local':  ctrl_conf_local[:,  img_id],
                    'track_pts3d_local': [pts3d_t_l[[t]] for t in range(num_images)],
                    'track_conf_local':  [conf_t_l[[t]]  for t in range(num_images)],
                })
            results[img_id] = res

        return results

	# ruff: noqa
	# fmt: off
	"""Trace Anything, vendored verbatim from ByteDance-Seed/TraceAnything @ 54677b5e7bf11510c2e8c917a509988ad379f8eb (trace_anything/): the files below
	concatenated, package-internal imports replaced by relative imports of the other vendored files (generated by scripts/model_ports/common/vendor_flat.py;
	do not edit by hand). Code license: Apache-2.0."""
	from .ta_heads import PixelHead, ScalarHead
	from .ta_layers import Block, PositionGetter, RoPE2D, get_1d_sincos_pos_embed_from_grid, get_patch_embed


	# ==============================================================================================================
	# trace_anything/trace_anything.py
	# ==============================================================================================================
	# Copyright (c) 2025 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.

	# trace_anything/trace_anything.py
	"""
	Minimal inference-only model for this repo.

	Assumptions (pruned by asserts):
	- encoder_type == 'croco'
	- decoder_type == 'transformer'
	- head_type == 'dpt'
	- targeting_mechanism == 'bspline_conf'
	- optional: whether_local (bool)
	"""

	import math
	import time
	from copy import deepcopy
	from functools import partial
	from typing import Dict, List

	import torch
	import torch.nn as nn
	from einops import rearrange
	import numpy as np


	from contextlib import contextmanager
	import time



	# ======== B-spline ========
	PRECOMPUTED_KNOTS = {
	4: torch.tensor([0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0], dtype=torch.float32),
	7: torch.tensor([0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.5, 1.0, 1.0, 1.0, 1.0], dtype=torch.float32),
	10: torch.tensor([0.0, 0.0, 0.0, 0.0, 1/3, 1/3, 1/3, 2/3, 2/3, 2/3, 1.0, 1.0, 1.0, 1.0], dtype=torch.float32),
	}

	def _precompute_knot_differences(n_ctrl_pts, degree, knots):
	denom1 = torch.zeros(n_ctrl_pts, degree + 1, device=knots.device)
	denom2 = torch.zeros(n_ctrl_pts, degree + 1, device=knots.device)
	for k in range(degree + 1):
	for i in range(n_ctrl_pts):
	denom1[i, k] = knots[i + k] - knots[i] if i + k < len(knots) else 0.0
	denom2[i, k] = knots[i + k + 1] - knots[i + 1] if i + k + 1 < len(knots) else 1.0
	return denom1, denom2

	PRECOMPUTED_DENOMS = {n: _precompute_knot_differences(n, 3, PRECOMPUTED_KNOTS[n]) for n in [4, 7, 10]}

	def _compute_bspline_basis(n_ctrl_pts, degree, t_values, knots, denom1, denom2):
	N = t_values.size(0)
	basis = torch.zeros(N, n_ctrl_pts, degree + 1, device=t_values.device)
	t = t_values
	basis_k0 = torch.zeros(N, n_ctrl_pts, device=t.device)
	for i in range(n_ctrl_pts):
	if i == n_ctrl_pts - 1:
	basis_k0[:, i] = ((knots[i] <= t) & (t <= knots[i + 1])).float()
	else:
	basis_k0[:, i] = ((knots[i] <= t) & (t < knots[i + 1])).float()
	basis[:, :, 0] = basis_k0
	for k in range(1, degree + 1):
	basis_k = torch.zeros(N, n_ctrl_pts, device=t.device)
	for i in range(n_ctrl_pts):
	term1 = ((t - knots[i]) / denom1[i, k]) * basis[:, i, k-1] if denom1[i, k] > 0 else 0.0
	term2 = ((knots[i + k + 1] - t) / denom2[i, k]) * basis[:, i + 1, k-1] if (denom2[i, k] > 0 and i + 1 < n_ctrl_pts) else 0.0
	basis_k[:, i] = term1 + term2
	basis[:, :, k] = basis_k
	return basis[:, :, degree]

	def evaluate_bspline_conf(ctrl_pts3d, ctrl_conf, t_values):
	"""ctrl_pts3d:[N_ctrl,H,W,3], ctrl_conf:[N_ctrl,H,W], t_values:[T] -> (T,H,W,3),(T,H,W)"""
	n_ctrl_pts, H, W, _ = ctrl_pts3d.shape
	assert n_ctrl_pts in (4, 7, 10), f"unsupported n_ctrl_pts={n_ctrl_pts}"
	degree = 3
	knot_vector = PRECOMPUTED_KNOTS[n_ctrl_pts].to(ctrl_pts3d.device)
	denom1, denom2 = [d.to(ctrl_pts3d.device) for d in PRECOMPUTED_DENOMS[n_ctrl_pts]]
	ctrl_pts3d = ctrl_pts3d.permute(0, 3, 1, 2) # [N,3,H,W]
	ctrl_conf = ctrl_conf.unsqueeze(-1).permute(0, 3, 1, 2) # [N,1,H,W]
	basis = _compute_bspline_basis(n_ctrl_pts, degree, t_values, knot_vector, denom1, denom2) # [T,N]
	basis = basis.view(-1, n_ctrl_pts, 1, 1, 1) # [T,N,1,1,1]
	pts3d_t = torch.sum(basis * ctrl_pts3d.unsqueeze(0), dim=1).permute(0, 2, 3, 1) # [T,H,W,3]
	conf_t = torch.sum(basis * ctrl_conf.unsqueeze(0), dim=1).squeeze(1) # [T,H,W]
	return pts3d_t, conf_t

	# ======== Encoders（仅 CroCo） ========
	class CroCoEncoder(nn.Module):
	def __init__(
	self,
	img_size=512, patch_size=16, patch_embed_cls="ManyAR_PatchEmbed",
	embed_dim=768, num_heads=12, depth=12, mlp_ratio=4,
	norm_layer=partial(nn.LayerNorm, eps=1e-6),
	pos_embed="RoPE100", attn_implementation="pytorch_naive",
	):
	super().__init__()
	assert pos_embed.startswith("RoPE"), f"pos_embed must start with RoPE*, got {pos_embed}"
	self.patch_embed = get_patch_embed(patch_embed_cls, img_size, patch_size, embed_dim)
	freq = float(pos_embed[len("RoPE"):])
	self.rope = RoPE2D(freq=freq)
	self.enc_blocks = nn.ModuleList([
	Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio,
	qkv_bias=True, norm_layer=nn.LayerNorm, rope=self.rope,
	attn_implementation=attn_implementation)
	for _ in range(depth)
	])
	self.enc_norm = norm_layer(embed_dim)

	def forward(self, image, true_shape):
	x, pos = self.patch_embed(image, true_shape=true_shape)
	for blk in self.enc_blocks:
	x = blk(x, pos)
	x = self.enc_norm(x)
	return x, pos


	# ======== Decoder ========
	class TraceDecoder(nn.Module):
	def __init__(
	self,
	random_image_idx_embedding: bool,
	enc_embed_dim: int,
	embed_dim: int = 768,
	num_heads: int = 12,
	depth: int = 12,
	mlp_ratio: float = 4.0,
	qkv_bias: bool = True,
	drop: float = 0.0,
	attn_drop: float = 0.0,
	attn_implementation: str = "pytorch_naive",
	attn_bias_for_inference_enabled=True,
	norm_layer=partial(nn.LayerNorm, eps=1e-6),
	):
	super().__init__()
	self.decoder_embed = nn.Linear(enc_embed_dim, embed_dim, bias=True)
	self.dec_blocks = nn.ModuleList([
	Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio,
	qkv_bias=qkv_bias, drop=drop, attn_drop=attn_drop,
	norm_layer=nn.LayerNorm, attn_implementation=attn_implementation,
	attn_bias_for_inference_enabled=attn_bias_for_inference_enabled)
	for _ in range(depth)
	])
	self.random_image_idx_embedding = random_image_idx_embedding
	self.register_buffer(
	"image_idx_emb",
	torch.from_numpy(get_1d_sincos_pos_embed_from_grid(embed_dim, np.arange(1000))).float(),
	persistent=False,
	)
	self.dec_norm = norm_layer(embed_dim)

	def _get_random_image_pos(self, encoded_feats, batch_size, num_views, max_image_idx, device):
	image_ids = torch.zeros(batch_size, num_views, dtype=torch.long)
	image_ids[:, 0] = 0
	per_forward_pass_seed = torch.randint(0, 2 ** 32, (1,)).item()
	per_rank_generator = torch.Generator().manual_seed(per_forward_pass_seed)
	for b in range(batch_size):
	random_ids = torch.randperm(max_image_idx, generator=per_rank_generator)[:num_views - 1] + 1
	image_ids[b, 1:] = random_ids
	image_ids = image_ids.to(device)
	image_pos_list = []
	for i in range(num_views):
	num_patches = encoded_feats[i].shape[1]
	pos_for_view = self.image_idx_emb[image_ids[:, i]].unsqueeze(1).repeat(1, num_patches, 1)
	image_pos_list.append(pos_for_view)
	return torch.cat(image_pos_list, dim=1)

	def forward(self, encoded_feats: List[torch.Tensor], positions: List[torch.Tensor],
	image_ids: torch.Tensor, image_timesteps: torch.Tensor):
	x = torch.cat(encoded_feats, dim=1)
	pos = torch.cat(positions, dim=1)
	outputs = [x]
	x = self.decoder_embed(x)

	if self.random_image_idx_embedding:
	image_pos = self._get_random_image_pos(
	encoded_feats=encoded_feats,
	batch_size=encoded_feats[0].shape[0],
	num_views=len(encoded_feats),
	max_image_idx=self.image_idx_emb.shape[0] - 1,
	device=x.device,
	)
	else:
	num_embeddings = self.image_idx_emb.shape[0]
	indices = (image_timesteps * (num_embeddings - 1)).long().view(-1)
	image_pos = torch.index_select(self.image_idx_emb, dim=0, index=indices)
	image_pos = image_pos.view(1, image_timesteps.shape[1], self.image_idx_emb.shape[1])

	x += image_pos
	for blk in self.dec_blocks:
	x = blk(x, pos)
	outputs.append(x)

	x = self.dec_norm(x)
	outputs[-1] = x
	return outputs, image_pos


	# ======== main model ========
	class TraceAnything(nn.Module):
	def __init__(self, *, encoder_args: Dict, decoder_args: Dict, head_args: Dict,
	targeting_mechanism: str = "bspline_conf",
	poly_degree: int = 10, whether_local: bool = False):
	super().__init__()

	assert targeting_mechanism == "bspline_conf", f"Only bspline_conf is supported now, got {targeting_mechanism}"
	assert encoder_args.get("encoder_type") == "croco"
	assert decoder_args.get("decoder_type", "transformer") in ("transformer", "fast3r")
	assert head_args.get("head_type") == "dpt"

	self.targeting_mechanism = targeting_mechanism
	self.poly_degree = int(poly_degree)
	self.whether_local = bool(whether_local or head_args.get("with_local_head", False))

	# build encoder / decoder
	enc_args = deepcopy(encoder_args); enc_args.pop("encoder_type", None)
	self.encoder = CroCoEncoder(**enc_args)

	dec_args = deepcopy(decoder_args); dec_args.pop("decoder_type", None)
	self.decoder = TraceDecoder(**dec_args)

	# build heads
	feature_dim = 256
	last_dim = feature_dim // 2
	ed = encoder_args["embed_dim"]; dd = decoder_args["embed_dim"]
	hooks = [0, decoder_args["depth"] * 2 // 4, decoder_args["depth"] * 3 // 4, decoder_args["depth"]]
	self.ds_head_time = ScalarHead(
	num_channels=1, feature_dim=feature_dim, last_dim=last_dim,
	hooks_idx=hooks, dim_tokens=[ed, dd, dd, dd], head_type="regression",
	patch_size=head_args["patch_size"],
	)
	out_nchan = (3 + bool(head_args["conf_mode"])) * self.poly_degree
	self.ds_head_track = PixelHead(
	num_channels=out_nchan, feature_dim=feature_dim, last_dim=last_dim,
	hooks_idx=hooks, dim_tokens=[ed, dd, dd, dd], head_type="regression",
	patch_size=head_args["patch_size"],
	depth_mode=head_args["depth_mode"], conf_mode=head_args["conf_mode"],
	)
	if self.whether_local:
	self.ds_head_local = PixelHead(
	num_channels=out_nchan, feature_dim=feature_dim, last_dim=last_dim,
	hooks_idx=hooks, dim_tokens=[ed, dd, dd, dd], head_type="regression",
	patch_size=head_args["patch_size"],
	depth_mode=head_args["depth_mode"], conf_mode=head_args["conf_mode"],
	)

	self.time_head = self.ds_head_time
	self.track_head = self.ds_head_track
	if self.whether_local:
	self.local_head = self.ds_head_local

	self.max_parallel_views_for_head = 25

	def _encode_images(self, views, chunk_size=400):
	B = views[0]["img"].shape[0]
	same_shape = all(v["img"].shape == views[0]["img"].shape for v in views)
	if same_shape:
	imgs = torch.cat([v["img"] for v in views], dim=0)
	true_shapes = torch.cat([v.get("true_shape", torch.tensor(v["img"].shape[-2:])[None].repeat(B, 1)) for v in views], dim=0)
	feats_chunks, pos_chunks = [], []
	for s in range(0, imgs.shape[0], chunk_size):
	e = min(s + chunk_size, imgs.shape[0])
	f, p = self.encoder(imgs[s:e], true_shapes[s:e])
	feats_chunks.append(f); pos_chunks.append(p)
	feats = torch.cat(feats_chunks, dim=0); pos = torch.cat(pos_chunks, dim=0)
	encoded_feats = torch.split(feats, B, dim=0)
	positions = torch.split(pos, B, dim=0)
	shapes = torch.split(true_shapes, B, dim=0)
	else:
	encoded_feats, positions, shapes = [], [], []
	for v in views:
	img = v["img"]
	true_shape = v.get("true_shape", torch.tensor(img.shape[-2:])[None].repeat(B, 1))
	f, p = self.encoder(img, true_shape)
	encoded_feats.append(f); positions.append(p); shapes.append(true_shape)
	return encoded_feats, positions, shapes

	@torch.no_grad()
	def forward(self, views, profiling: bool = False):
	# 1) encode
	encoded_feats, positions, shapes = self._encode_images(views)

	# 2) build time embedding
	num_images = len(views)
	B, P, D = encoded_feats[0].shape
	image_ids, image_times = [], []
	for i, ef in enumerate(encoded_feats):
	num_patches = ef.shape[1]
	image_ids.extend([i] * num_patches)
	image_times.extend([views[i]["time_step"]] * num_patches)
	image_ids = torch.tensor(image_ids * B, device=encoded_feats[0].device).reshape(B, -1)
	image_times = torch.tensor(image_times * B, device=encoded_feats[0].device).reshape(B, -1)

	# 3) decode
	dec_output, _ = self.decoder(encoded_feats, positions, image_ids, image_times)

	# 4) gather outputs per view
	P_patches = P
	gathered_outputs_list = []
	for layer_output in dec_output:
	layer_output = rearrange(layer_output, 'B (n P) D -> (n B) P D', n=num_images, P=P_patches)
	gathered_outputs_list.append(layer_output)

	# 5) heads
	time_step = self.time_head(gathered_outputs_list, torch.stack(shapes)) # [N, 1] per view
	track_tmp = self.track_head(gathered_outputs_list, torch.stack(shapes)) # dict with pts3d/conf after postprocess
	ctrl_pts3d, ctrl_conf = track_tmp['pts3d'], track_tmp['conf'] # [K,N,H,W,3], [K,N,H,W]

	if self.whether_local:
	local_tmp = self.local_head(gathered_outputs_list, torch.stack(shapes))
	ctrl_pts3d_local, ctrl_conf_local = local_tmp['pts3d'], local_tmp['conf']

	# 6) evaluate bspline track over all reference times
	results = [{} for _ in range(num_images)]
	t_values = torch.stack([time_step[i].squeeze() for i in range(num_images)]) # [N]
	for img_id in range(num_images):
	pts3d_t, conf_t = evaluate_bspline_conf(ctrl_pts3d[:, img_id], ctrl_conf[:, img_id], t_values.detach())
	res = {
	'time': time_step[img_id],
	'ctrl_pts3d': ctrl_pts3d[:, img_id],
	'ctrl_conf': ctrl_conf[:, img_id],
	'track_pts3d': [pts3d_t[[t]] for t in range(num_images)],
	'track_conf': [conf_t[[t]] for t in range(num_images)],
	}
	if self.whether_local:
	pts3d_t_l, conf_t_l = evaluate_bspline_conf(ctrl_pts3d_local[:, img_id], ctrl_conf_local[:, img_id], t_values.detach())
	res.update({
	'ctrl_pts3d_local': ctrl_pts3d_local[:, img_id],
	'ctrl_conf_local': ctrl_conf_local[:, img_id],
	'track_pts3d_local': [pts3d_t_l[[t]] for t in range(num_images)],
	'track_conf_local': [conf_t_l[[t]] for t in range(num_images)],
	})
	results[img_id] = res

	return results