mstar.model.waypoint.components.taehv

Functions

decode_latent(ae_model, latent, histories, ...)	Pure decoder initialization or steady-state step.
decoder_history_shapes(ae_model, *, ...)	The nine fixed MemBlock input shapes for one decoder request.
encode_seed_clip(ae_model, frames, *, ...)	Encode one fixed four-frame RGB clip into one latent tensor.
encoded_size_for_latent(latent_height, ...)
initial_decoder_histories(ae_model, latent)
load_taehv(ae_uri[, cache_dir])	Load the shared weights module from a resolved local checkpoint.
pixel_size_for_latent(latent_height, ...)
validate_taehv_architecture(ae_model)	Validate the fixed architecture the captured functional path supports.

Classes

ChunkedStreamingTAEHV(ae_model[, ...])

One streaming session: this request's frames, in order.

class mstar.model.waypoint.components.taehv.ChunkedStreamingTAEHV(ae_model, auto_aspect_ratio=True, device=None, dtype=torch.bfloat16, height=None, width=None)

Bases: object

One streaming session: this request’s frames, in order.

The order-dependent state lives here; ae_model is shared and never mutated, so a fresh instance is the reference’s reset(). Not an nn.Module: the module tree would give state_dict(), to() and the weight loader a second path to the same parameters.

Parameters:

• ae_model (Module)

• auto_aspect_ratio (bool)

• device (device | None)

• dtype (dtype)

• height (int | None)

• width (int | None)

decode(latent)

One latent [B, C, h, w] -> [t_upscale, H, W, 3] uint8.

Single-use and order-dependent: the temporal memory advances per call, so a latent decoded twice, out of turn or not at all shifts every frame after it, silently. The first call primes that memory with frames_to_trim extra feeds, which is why the prime walk decodes.

Parameters:

latent (Tensor)

Return type:

Tensor

encode(frames)

[t_downscale, H, W, 3] in [0, 1] -> one latent [B, C, h, w]. Exactly that many frames: fewer buffer and return None. The reference scales from uint8 here; this takes it pre-scaled in the same cast-then-divide order, so the two agree bit for bit.

Parameters:

frames (Tensor)

Return type:

Tensor

mstar.model.waypoint.components.taehv.decode_latent(ae_model, latent, histories, *, output_size, initialize)

Pure decoder initialization or steady-state step.

Initialization reproduces ChunkedStreamingTAEHV.decode exactly: the seed latent is fed frames_to_trim extra times, those reconstructed clips are discarded, and a final feed advances the state through the seed frame. The caller intentionally emits none of those frames.

Parameters:

• ae_model (Module)

• latent (Tensor)

• histories (Sequence[Tensor])

• output_size (tuple[int, int])

• initialize (bool)

Return type:

tuple[Tensor, tuple[Tensor, …]]

mstar.model.waypoint.components.taehv.decoder_history_shapes(ae_model, *, batch_size, latent_height, latent_width)

The nine fixed MemBlock input shapes for one decoder request.

Parameters:

• ae_model (Module)

• batch_size (int)

• latent_height (int)

• latent_width (int)

Return type:

tuple[tuple[int, int, int, int], …]

mstar.model.waypoint.components.taehv.encode_seed_clip(ae_model, frames, *, output_size)

Encode one fixed four-frame RGB clip into one latent tensor.

The operation is pure with respect to ae_model: only its parameters are read, and every temporal history is created and consumed within this call.

Parameters:

• ae_model (Module)

• frames (Tensor)

• output_size (tuple[int, int])

Return type:

Tensor

mstar.model.waypoint.components.taehv.encoded_size_for_latent(latent_height, latent_width)

Parameters:

• latent_height (int)

• latent_width (int)

Return type:

tuple[int, int]

mstar.model.waypoint.components.taehv.initial_decoder_histories(ae_model, latent)

Parameters:

• ae_model (Module)

• latent (Tensor)

Return type:

tuple[Tensor, …]

mstar.model.waypoint.components.taehv.load_taehv(ae_uri, cache_dir=None)

Load the shared weights module from a resolved local checkpoint.

Hub resolution belongs to checkpoint.resolve_taehv_checkpoint, so only the required file is downloaded. cache_dir is accepted for compatibility with existing callers but is unused here.

Parameters:

• ae_uri (str)

• cache_dir (str | None)

Return type:

Module

mstar.model.waypoint.components.taehv.pixel_size_for_latent(latent_height, latent_width)

Parameters:

• latent_height (int)

• latent_width (int)

Return type:

tuple[int, int]

mstar.model.waypoint.components.taehv.validate_taehv_architecture(ae_model)

Validate the fixed architecture the captured functional path supports.

Parameters:

ae_model (Module)

Return type:

None