Trace-Anything / modeling_trace_anything.py

Upload Trace-Anything (transformers port, original code)

653a5eb verified 14 days ago

5.99 kB

"""Trace Anything (ByteDance-Seed/TraceAnything @ 54677b5e) as a transformers model.
``TraceAnythingModel`` holds the submodules of the official ``TraceAnything`` under the same names (``encoder``: CroCo
ViT, ``decoder``: the trace decoder, ``ds_head_time`` / ``ds_head_track``: DPT heads; the official ``time_head`` /
``track_head`` are aliases of the latter two), so the released ``trace_anything.pt`` tensors load by name, and runs the
official ``TraceAnything.forward`` on the views built from ``pixel_values`` / ``time_steps``. The per-view result dicts
of the original are returned as ``predictions`` and stacked over the views in ``TraceAnythingOutput``.
"""

from dataclasses import dataclass

import numpy as np
import torch
from transformers import PreTrainedModel
from transformers.utils import ModelOutput

from .configuration_trace_anything import TraceAnythingConfig
from .ta_layers import get_1d_sincos_pos_embed_from_grid
from .ta_model import TraceAnything, TraceDecoder


@dataclass
class TraceAnythingOutput(ModelOutput):
    """Per view i of N (stacked from the original's per-view dicts):
    time (N, 1): the predicted time of each view; ctrl_pts3d (N, K, H, W, 3) / ctrl_conf (N, K, H, W): the K B-spline
    control points of each pixel's 3D trajectory and their confidence; track_pts3d (N, N, H, W, 3) / track_conf
    (N, N, H, W): the trajectories evaluated at the N predicted times; predictions: the original list of dicts."""

    time: torch.Tensor | None = None
    ctrl_pts3d: torch.Tensor | None = None
    ctrl_conf: torch.Tensor | None = None
    track_pts3d: torch.Tensor | None = None
    track_conf: torch.Tensor | None = None
    predictions: list | None = None


class TraceAnythingPreTrainedModel(PreTrainedModel):
    config_class = TraceAnythingConfig
    config: TraceAnythingConfig
    base_model_prefix = "trace_anything"
    main_input_name = "pixel_values"

    def _init_weights(self, module):
        super()._init_weights(module)
        # TraceDecoder.image_idx_emb is a non-persistent buffer: not in checkpoints, left uninitialised by meta-device
        # loading -> rebuild it as TraceDecoder.__init__ does
        if isinstance(module, TraceDecoder):
            embed_dim = module.image_idx_emb.shape[1]
            with torch.no_grad():
                module.image_idx_emb.copy_(torch.from_numpy(get_1d_sincos_pos_embed_from_grid(embed_dim, np.arange(1000))).float())


class TraceAnythingModel(TraceAnythingPreTrainedModel):
    _encode_images = TraceAnything._encode_images

    def __init__(self, config: TraceAnythingConfig):
        super().__init__(config)
        net = TraceAnything(encoder_args=config.encoder_args, decoder_args=config.decoder_args, head_args=config.head_args,
                            targeting_mechanism=config.targeting_mechanism, poly_degree=config.poly_degree, whether_local=False)
        if net.whether_local:
            raise ValueError("head_args.with_local_head must be False: no released checkpoint has local-head weights")
        # the official module's attributes, so its methods run on this object unchanged
        self.targeting_mechanism = net.targeting_mechanism
        self.poly_degree = net.poly_degree
        self.whether_local = net.whether_local
        self.max_parallel_views_for_head = net.max_parallel_views_for_head
        self.encoder = net.encoder
        self.decoder = net.decoder
        self.ds_head_time = net.ds_head_time
        self.ds_head_track = net.ds_head_track
        # the DPT heads register their scratch convolutions twice (layer1_rn ... and the layer_rn list): declare those
        # names as tied (target -> first name), so save_pretrained stores each tensor once and from_pretrained keeps the sharing
        first, tied = {}, {}
        for name, parameter in self.named_parameters(remove_duplicate=False):
            if id(parameter) in first:
                tied[name] = first[id(parameter)]
            else:
                first[id(parameter)] = name
        self._tied_weights_keys = tied or None
        self.post_init()

    # the official aliases (self.time_head = self.ds_head_time, ...), not registered twice
    @property
    def time_head(self):
        return self.ds_head_time

    @property
    def track_head(self):
        return self.ds_head_track

    def forward(self, pixel_values: torch.Tensor, time_steps=None, return_dict: bool | None = None, **kwargs) -> tuple | TraceAnythingOutput:
        """pixel_values (1, N, 3, H, W) in [-1, 1], H and W multiples of 16; time_steps (1, N): the time of each frame
        (default i / (N - 1)); both from ``TraceAnythingProcessor``. One video per call, as the original."""
        return_dict = return_dict if return_dict is not None else self.config.return_dict
        if pixel_values.shape[0] != 1:
            raise ValueError("one video per call: pixel_values must be (1, N, 3, H, W)")
        num_views = pixel_values.shape[1]
        if time_steps is None:
            times = [i / (num_views - 1) if num_views > 1 else 0.0 for i in range(num_views)]
        else:
            times = [float(t) for t in (time_steps[0].tolist() if isinstance(time_steps, torch.Tensor) else time_steps[0])]
        views = [{"img": pixel_values[:, i], "time_step": times[i]} for i in range(num_views)]
        predictions = TraceAnything.forward(self, views)
        output = TraceAnythingOutput(
            time=torch.stack([p["time"] for p in predictions]),
            ctrl_pts3d=torch.stack([p["ctrl_pts3d"] for p in predictions]),
            ctrl_conf=torch.stack([p["ctrl_conf"] for p in predictions]),
            track_pts3d=torch.stack([torch.cat(p["track_pts3d"]) for p in predictions]),
            track_conf=torch.stack([torch.cat(p["track_conf"]) for p in predictions]),
            predictions=predictions,
        )
        return output if return_dict else output.to_tuple()


__all__ = ["TraceAnythingModel", "TraceAnythingOutput", "TraceAnythingPreTrainedModel"]

	"""Trace Anything (ByteDance-Seed/TraceAnything @ 54677b5e) as a transformers model.

	``TraceAnythingModel`` holds the submodules of the official ``TraceAnything`` under the same names (``encoder``: CroCo
	ViT, ``decoder``: the trace decoder, ``ds_head_time`` / ``ds_head_track``: DPT heads; the official ``time_head`` /
	``track_head`` are aliases of the latter two), so the released ``trace_anything.pt`` tensors load by name, and runs the
	official ``TraceAnything.forward`` on the views built from ``pixel_values`` / ``time_steps``. The per-view result dicts
	of the original are returned as ``predictions`` and stacked over the views in ``TraceAnythingOutput``.
	"""

	from dataclasses import dataclass

	import numpy as np
	import torch
	from transformers import PreTrainedModel
	from transformers.utils import ModelOutput

	from .configuration_trace_anything import TraceAnythingConfig
	from .ta_layers import get_1d_sincos_pos_embed_from_grid
	from .ta_model import TraceAnything, TraceDecoder


	@dataclass
	class TraceAnythingOutput(ModelOutput):
	"""Per view i of N (stacked from the original's per-view dicts):
	time (N, 1): the predicted time of each view; ctrl_pts3d (N, K, H, W, 3) / ctrl_conf (N, K, H, W): the K B-spline
	control points of each pixel's 3D trajectory and their confidence; track_pts3d (N, N, H, W, 3) / track_conf
	(N, N, H, W): the trajectories evaluated at the N predicted times; predictions: the original list of dicts."""

	time: torch.Tensor \| None = None
	ctrl_pts3d: torch.Tensor \| None = None
	ctrl_conf: torch.Tensor \| None = None
	track_pts3d: torch.Tensor \| None = None
	track_conf: torch.Tensor \| None = None
	predictions: list \| None = None


	class TraceAnythingPreTrainedModel(PreTrainedModel):
	config_class = TraceAnythingConfig
	config: TraceAnythingConfig
	base_model_prefix = "trace_anything"
	main_input_name = "pixel_values"

	def _init_weights(self, module):
	super()._init_weights(module)
	# TraceDecoder.image_idx_emb is a non-persistent buffer: not in checkpoints, left uninitialised by meta-device
	# loading -> rebuild it as TraceDecoder.__init__ does
	if isinstance(module, TraceDecoder):
	embed_dim = module.image_idx_emb.shape[1]
	with torch.no_grad():
	module.image_idx_emb.copy_(torch.from_numpy(get_1d_sincos_pos_embed_from_grid(embed_dim, np.arange(1000))).float())


	class TraceAnythingModel(TraceAnythingPreTrainedModel):
	_encode_images = TraceAnything._encode_images

	def __init__(self, config: TraceAnythingConfig):
	super().__init__(config)
	net = TraceAnything(encoder_args=config.encoder_args, decoder_args=config.decoder_args, head_args=config.head_args,
	targeting_mechanism=config.targeting_mechanism, poly_degree=config.poly_degree, whether_local=False)
	if net.whether_local:
	raise ValueError("head_args.with_local_head must be False: no released checkpoint has local-head weights")
	# the official module's attributes, so its methods run on this object unchanged
	self.targeting_mechanism = net.targeting_mechanism
	self.poly_degree = net.poly_degree
	self.whether_local = net.whether_local
	self.max_parallel_views_for_head = net.max_parallel_views_for_head
	self.encoder = net.encoder
	self.decoder = net.decoder
	self.ds_head_time = net.ds_head_time
	self.ds_head_track = net.ds_head_track
	# the DPT heads register their scratch convolutions twice (layer1_rn ... and the layer_rn list): declare those
	# names as tied (target -> first name), so save_pretrained stores each tensor once and from_pretrained keeps the sharing
	first, tied = {}, {}
	for name, parameter in self.named_parameters(remove_duplicate=False):
	if id(parameter) in first:
	tied[name] = first[id(parameter)]
	else:
	first[id(parameter)] = name
	self._tied_weights_keys = tied or None
	self.post_init()

	# the official aliases (self.time_head = self.ds_head_time, ...), not registered twice
	@property
	def time_head(self):
	return self.ds_head_time

	@property
	def track_head(self):
	return self.ds_head_track

	def forward(self, pixel_values: torch.Tensor, time_steps=None, return_dict: bool \| None = None, **kwargs) -> tuple \| TraceAnythingOutput:
	"""pixel_values (1, N, 3, H, W) in [-1, 1], H and W multiples of 16; time_steps (1, N): the time of each frame
	(default i / (N - 1)); both from ``TraceAnythingProcessor``. One video per call, as the original."""
	return_dict = return_dict if return_dict is not None else self.config.return_dict
	if pixel_values.shape[0] != 1:
	raise ValueError("one video per call: pixel_values must be (1, N, 3, H, W)")
	num_views = pixel_values.shape[1]
	if time_steps is None:
	times = [i / (num_views - 1) if num_views > 1 else 0.0 for i in range(num_views)]
	else:
	times = [float(t) for t in (time_steps[0].tolist() if isinstance(time_steps, torch.Tensor) else time_steps[0])]
	views = [{"img": pixel_values[:, i], "time_step": times[i]} for i in range(num_views)]
	predictions = TraceAnything.forward(self, views)
	output = TraceAnythingOutput(
	time=torch.stack([p["time"] for p in predictions]),
	ctrl_pts3d=torch.stack([p["ctrl_pts3d"] for p in predictions]),
	ctrl_conf=torch.stack([p["ctrl_conf"] for p in predictions]),
	track_pts3d=torch.stack([torch.cat(p["track_pts3d"]) for p in predictions]),
	track_conf=torch.stack([torch.cat(p["track_conf"]) for p in predictions]),
	predictions=predictions,
	)
	return output if return_dict else output.to_tuple()


	__all__ = ["TraceAnythingModel", "TraceAnythingOutput", "TraceAnythingPreTrainedModel"]