minimax_h3_video_vae ¶

def decode(self,
           z,
           *,
           tiled: bool = True,
           tile_sample_min_height: int = 256,
           tile_sample_min_width: int = 256,
           tile_sample_min_overlap_height: int = 64,
           tile_sample_min_overlap_width: int = 64):
    """Chunked decode of normalized latents (1, C, T_lat, H', W') -> (1, 3, T, H, W)."""
    cfg = self.config
    tokens_chunk_size = cfg.tokens_chunk_size
    token_drop = cfg.token_drop
    temporal_ratio = self.temporal_compression_ratio
    chunk_num_frames = tokens_chunk_size * temporal_ratio
    num_tokens = z.shape[2] + token_drop
    pad_tokens = (-num_tokens) % tokens_chunk_size
    num_chunks = (num_tokens + pad_tokens) // tokens_chunk_size - int(token_drop > 0)
    if pad_tokens > 0:
        tail = mx.repeat(z[:, :, -1:], pad_tokens, axis=2)
        z = mx.concatenate([z, tail], axis=2)

    def decode_one(chunk):
        if tiled:
            return self.decode_clip_tiled(chunk, tile_sample_min_height, tile_sample_min_width,
                                          tile_sample_min_overlap_height, tile_sample_min_overlap_width)
        return self._decode_clip(chunk)

    decoded_chunks = []
    overlap = None
    for index in range(num_chunks):
        start = index * tokens_chunk_size
        clip = decode_one(z[:, :, start:start + tokens_chunk_size + cfg.token_overlap])
        for overlap_index in range(int(token_drop > 0) + 1):
            frame_start = overlap_index * chunk_num_frames
            chunk = clip[:, :, frame_start:frame_start + chunk_num_frames]
            chunk = chunk[:, :, cfg.frame_pre_padding:]
            if overlap_index == 0:
                if overlap is not None:
                    chunk = self._blend(overlap, chunk, cfg.frame_overlap, dim=-3)
                mx.eval(chunk)  # materialize per clip; keeps the lazy graph bounded
                decoded_chunks.append(chunk)
            else:
                overlap = chunk
                mx.eval(overlap)
    if overlap is not None:
        decoded_chunks.append(overlap)
    decoded = mx.concatenate(decoded_chunks, axis=2)

    if pad_tokens > 0:
        intra_tail = cfg.clip_length % temporal_ratio
        num_tokens_before_pad = z.shape[2] - pad_tokens
        pad_frames = sum(intra_tail if intra_tail and (num_tokens_before_pad + offset) %
                         tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens))
        decoded = decoded[:, :, :-pad_frames]
    return decoded

def decode_clip_tiled(self,
                      z,
                      tile_sample_min_height: int,
                      tile_sample_min_width: int,
                      tile_sample_min_overlap_height: int = 64,
                      tile_sample_min_overlap_width: int = 64):
    """One clip through the decoder with spatial tiling (memory bounded)."""
    height = z.shape[-2] * self.spatial_compression_ratio
    width = z.shape[-1] * self.spatial_compression_ratio
    y_starts, y_lengths, y_overlaps = self._split_tiles(height, tile_sample_min_height,
                                                        tile_sample_min_overlap_height)
    x_starts, x_lengths, x_overlaps = self._split_tiles(width, tile_sample_min_width, tile_sample_min_overlap_width)
    ratio = self.spatial_compression_ratio
    rows = []
    for y_start, y_length in zip(y_starts, y_lengths, strict=False):
        row = []
        for x_start, x_length in zip(x_starts, x_lengths, strict=False):
            tile = z[..., y_start // ratio:y_start // ratio + y_length // ratio,
                     x_start // ratio:x_start // ratio + x_length // ratio]
            row.append(self._decode_clip(tile))
        rows.append(row)
    return self._stitch_tiles(rows, y_overlaps, x_overlaps)

def encode(self, pixels):
    """Normalized pixels (1, 3, T, H, W) -> (mean, logvar) each (1, C, T', H', W').

    Pads the clip to ``clip_length`` and drops ``token_drop`` trailing
    moment frames exactly like the reference ``_encode``.
    """
    if not self.has_encoder:
        raise RuntimeError("This MLX H3 video VAE was loaded without encoder weights.")
    cfg = self.config
    num_frames = pixels.shape[2]
    if num_frames % cfg.clip_length != 0:
        pad_frames = (-num_frames) % cfg.clip_length
        tail = mx.repeat(pixels[:, :, -1:], pad_frames, axis=2)
        pixels = mx.concatenate([pixels, tail], axis=2)
    moment_chunks: list[Any] = []
    for index in range(pixels.shape[2] // cfg.clip_length):
        clip = pixels[:, :, index * cfg.clip_length:(index + 1) * cfg.clip_length]
        moment_chunks.append(self._encode_clip(clip))
    moments = mx.concatenate(moment_chunks, axis=2)
    if cfg.token_drop > 0:
        moments = moments[:, :, :-cfg.token_drop]
    mean, logvar = mx.split(moments, 2, axis=1)
    logvar = mx.clip(logvar, -30.0, 20.0)
    return mean, logvar

def encode_keyframe(self, pixels):
    """Single-frame conditioning encode without chunk padding."""
    if pixels.shape[2] != 1:
        raise ValueError(f"encode_keyframe expects exactly one frame, got {pixels.shape}.")
    moments = self._encode_clip(pixels)
    mean, logvar = mx.split(moments, 2, axis=1)
    logvar = mx.clip(logvar, -30.0, 20.0)
    return mean, logvar

@staticmethod
def sample_posterior(mean, logvar, noise):
    """Reparameterization with an explicit noise array (deterministic parity)."""
    return mean + mx.exp(0.5 * logvar) * noise