FL2VA/video_vae/klvae.py

# SPDX-License-Identifier: Apache-2.0
# MiniMax H3 visual VAE: 3D causal CNN encoder + ViT3D decoder (inference-only bundle).
import os
import math
import numpy as np
import torch
import torch.nn as nn
import torch.distributed as dist
from typing import List, Union
from PIL import Image
from contextlib import nullcontext
from diffusers.models import ModelMixin
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.loaders.single_file_model import FromOriginalModelMixin
from diffusers.utils import logging

from .parallel import get_parallel_state, all_gather_var_shape
from .utils import apply_spatial_parallel
from .normalize import get_normalize_transform, get_denormalize_transform
from .vae_vit import ViT3DDecoder
from .vae_cnn import EncoderFCN3D
from .vae_module import DiagonalGaussianDistribution, ClsTokenAggregator
from .vae_processor import VAEProcessor


logger = logging.get_logger(__name__)  # pylint: disable=invalid-name


def _resolve_temporal_cat_dtype():
    raw = os.environ.get("MINIMAX_H3_VAE_DECODER_TEMPORAL_CAT_DTYPE", "").strip().lower()
    if raw in ("", "0", "false", "no", "off", "none", "keep", "default"):
        return None
    mapping = {
        "fp16": torch.float16,
        "float16": torch.float16,
        "half": torch.float16,
        "bf16": torch.bfloat16,
        "bfloat16": torch.bfloat16,
        "fp32": torch.float32,
        "float32": torch.float32,
    }
    if raw not in mapping:
        raise ValueError(
            "MINIMAX_H3_VAE_DECODER_TEMPORAL_CAT_DTYPE must be one of "
            "fp16|bf16|fp32|keep, got %r" % raw
        )
    return mapping[raw]


def _resolve_temporal_stream_cat():
    raw = os.environ.get("MINIMAX_H3_VAE_DECODER_STREAM_TEMPORAL_CAT", "1").strip().lower()
    return raw not in ("0", "false", "no", "off", "disable", "disabled")


class AutoencoderKL(ModelMixin, ConfigMixin, FromOriginalModelMixin):
    r"""
    Abstract shared base for the MiniMax H3 visual VAE.

    This class only carries the shared inference machinery (temporal
    chunking, tiling, encode/decode entry points). Instantiate the concrete
    subclass ``AutoencoderKLLegacy`` via ``from_pretrained`` instead.
    """

    _supports_gradient_checkpointing = True
    _compilable_modules = ["encoder", "decoder"]
    _deprecated_kwargs = [
        "clip_length",
        "token_drop",
        "isolated_first_frame",
        "isolated_last_frame",
        "isolated_key_frame",
        "encoder_tiling",
        "decoder_tiling",
        "parallel_tiling",
        "stack_tiling",
        "tile_size",
        "tile_overlap_min",
        "decoder_tile_size",
        "decoder_tile_overlap_min",
        "latent_patch_size",
        "crop_mode",
        "encoder_parallel",
        "decoder_parallel",
        "chunk_dim",
    ]  # legacy config keys accepted by from_pretrained for checkpoint compatibility


    def _set_gradient_checkpointing(self, module, value=False):
        if hasattr(module, "gradient_checkpointing"):
            module.gradient_checkpointing = value

    def _freeze_nested_module(self, module_path):
        parts = module_path.split(".")
        module = self
        for part in parts:
            module = getattr(module, part)
        module.requires_grad_(False)

    def setup_forward(self, **kwargs):
        self.clip_length = kwargs.get("clip_length", 17)
        self.token_drop = kwargs.get("token_drop", 0)
        self.frame_drop = self.token_drop * self.vae_ratio_t
        self.frame_pre_padding = (-self.clip_length) % self.vae_ratio_t
        self.tokens_chunk_size = math.ceil(self.clip_length / self.vae_ratio_t)
        self.token_overlap = (-self.token_drop) % self.tokens_chunk_size
        self.frame_overlap = max(self.token_overlap * self.vae_ratio_t - self.frame_pre_padding, 0)
        self.isolated_first_frame = kwargs.get("isolated_first_frame", False)
        self.isolated_last_frame = kwargs.get("isolated_last_frame", False)
        self.isolated_key_frame = kwargs.get("isolated_key_frame", False)

        self.encoder_tiling = kwargs.get("encoder_tiling", False)
        self.decoder_tiling = kwargs.get("decoder_tiling", False)
        self.stack_tiling = kwargs.get("stack_tiling", False)
        self.tile_size = kwargs.get("tile_size", 256)
        self.tile_overlap_min = kwargs.get("tile_overlap_min", 64)
        self.decoder_tile_size = kwargs.get("decoder_tile_size", self.tile_size)
        self.decoder_tile_overlap_min = kwargs.get("decoder_tile_overlap_min", self.tile_overlap_min)
        self.latent_patch_size = kwargs.get("latent_patch_size", 1)
        self.crop_mode = kwargs.get("crop_mode", "top_left")
        self.pixel_norm_type = kwargs.get("pixel_norm_type", "imagenet")

        # spatial parallel mode
        if hasattr(self, "_sp_initialized"):
            if (
                kwargs.get("chunk_dim", -1) != self.chunk_dim
                or kwargs.get("encoder_parallel", False) != self.encoder_parallel
                or kwargs.get("decoder_parallel", False) != self.decoder_parallel
                or kwargs.get("parallel_tiling", False) != self.parallel_tiling
            ):
                logger.warning(
                    "Do not support changing parallel schema after initialization"
                )
        else:
            self.chunk_dim = kwargs.get("chunk_dim", -1)
            self.encoder_parallel = kwargs.get("encoder_parallel", False)
            self.decoder_parallel = kwargs.get("decoder_parallel", False)
            self.parallel_tiling = kwargs.get("parallel_tiling", False)
            self._sp_initialized = True

        processor_kwargs = {
            "vae_ratio": self.vae_ratio,
            "vae_ratio_t": self.vae_ratio_t,
            "clip_length": self.clip_length,
            "frame_overlap": self.frame_overlap,
            "token_overlap": self.token_overlap,
            "tokens_chunk_size": self.tokens_chunk_size,
            "isolated_last_frame": self.isolated_last_frame,
            "latent_patch_size": self.latent_patch_size,
            "crop_mode": self.crop_mode,
            "pixel_norm_type": self.pixel_norm_type,
            "transform": self.transform,
            "transform_rev": self.transform_rev,
            "use_3d_conv": self.use_3d_conv,
        }
        if hasattr(self, "processor"):
            for key, value in processor_kwargs.items():
                setattr(self.processor, key, value)
        else:
            self.processor = VAEProcessor(**processor_kwargs)

    def perform_input_slice(self, x, chunk_size_stride=1):
        state = get_parallel_state()
        sp_rank = state["sp_rank"]
        sp_size = state["sp_size"]

        total_size = x.shape[self.chunk_dim]
        units = total_size // chunk_size_stride
        base_units = units // sp_size
        remainder_units = units % sp_size
        if sp_rank < remainder_units:
            start_units = sp_rank * (base_units + 1)
            end_units = start_units + base_units + 1
        else:
            start_units = sp_rank * base_units + remainder_units
            end_units = start_units + base_units
        start = start_units * chunk_size_stride
        end = end_units * chunk_size_stride

        slice_indices = [slice(None)] * x.ndim
        slice_indices[self.chunk_dim] = slice(start, end)
        x = x[tuple(slice_indices)].contiguous()
        return x

    def perform_output_concat(self, x):
        sp_process_group = get_parallel_state()["sp_process_group"]
        gathered = all_gather_var_shape(x, group=sp_process_group)
        x = torch.cat(gathered, dim=self.chunk_dim)
        return x



    def split_tiles(self, input_len, is_decoder=False):
        tile_size = self.decoder_tile_size if is_decoder else self.tile_size
        tile_overlap_min = self.decoder_tile_overlap_min if is_decoder else self.tile_overlap_min

        if tile_size >= input_len:
            return [0], [input_len], []

        N = math.ceil(input_len / tile_size)
        while True:
            overlaps = [tile_overlap_min] * (N - 1)
            remaining = tile_size * N - sum(overlaps) - input_len

            if remaining < 0:
                N += 1
            else:
                break

        remaining_units = remaining // self.vae_ratio
        for i in range(remaining_units):
            overlaps[i % (N - 1)] += self.vae_ratio

        tile_start_idx = [0]
        for i in range(N - 1):
            tile_start_idx.append(tile_start_idx[-1] + tile_size - overlaps[i])

        tile_len = [tile_size] * N
        return tile_start_idx, tile_len, overlaps

    def blend(
        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int, dim: int
    ) -> torch.Tensor:
        blend_extent = min(a.shape[dim], b.shape[dim], blend_extent)

        positions = torch.arange(blend_extent, device=b.device, dtype=b.dtype)
        weight_a = 1 - positions / blend_extent
        weight_b = positions / blend_extent

        shape = [1] * a.ndim
        shape[dim] = blend_extent
        weight_a = weight_a.view(shape)
        weight_b = weight_b.view(shape)

        slice_a = [slice(None)] * a.ndim
        slice_a[dim] = slice(-blend_extent, None)
        a_overlap = a[tuple(slice_a)]

        slice_b = [slice(None)] * b.ndim
        slice_b[dim] = slice(0, blend_extent)
        b_overlap = b[tuple(slice_b)]

        blended = a_overlap * weight_a + b_overlap * weight_b

        if blend_extent < b.shape[dim]:
            slice_b_rest = [slice(None)] * b.ndim
            slice_b_rest[dim] = slice(blend_extent, None)
            b_rest = b[tuple(slice_b_rest)]
            return torch.cat([blended, b_rest], dim=dim)
        else:
            return blended

    def _all_gather_tiled_results(self, tasks, num_tiles):
        state = get_parallel_state()
        group = state["sp_process_group"]
        sp_size = state["sp_size"]
        sp_rank = state["sp_rank"]

        if not tasks:
            raise ValueError(f"Found empty tasks on sp rank {sp_rank}")

        stacked = torch.stack(tasks, dim=0)
        gathered = all_gather_var_shape(stacked, group=group)

        results = [None] * num_tiles
        for rank, rank_tensors in enumerate(gathered):
            num_rank_tasks = rank_tensors.shape[0]
            for k in range(num_rank_tasks):
                global_idx = k * sp_size + rank
                if global_idx >= num_tiles:
                    break
                results[global_idx] = rank_tensors[k]

        return results

    def _local_tile_indices(self, num_tiles, sp_rank, sp_size):
        return list(range(sp_rank, num_tiles, sp_size))

    def _run_tile_tasks(self, tiles, tile_indices, forward_fn, stack_tiling, cls_agg=None):
        if stack_tiling and tile_indices:
            sample_batch_size = tiles[0].shape[0]
            tile_batch = torch.cat([tiles[idx] for idx in tile_indices], dim=0)
            output_batch = forward_fn(tile_batch)
            output_tiles = output_batch.unflatten(
                0, (len(tile_indices), sample_batch_size)
            ).unbind(dim=0)
            if cls_agg is not None:
                cls_agg.collect_stacked(len(tile_indices), sample_batch_size)
            return list(output_tiles)

        tasks = []
        for idx in tile_indices:
            tasks.append(forward_fn(tiles[idx]))
            if cls_agg is not None:
                cls_agg.collect()
        return tasks

    def tiled_encode(self, x):
        if self.parallel_tiling:  # Fast online encoding for large videos
            state = get_parallel_state()
            sp_rank = state["sp_rank"]
            sp_size = state["sp_size"]
        else:
            sp_rank, sp_size = 0, 1

        height, width = x.shape[-2], x.shape[-1]
        y_idx, y_len, y_overlap = self.split_tiles(height, False)
        x_idx, x_len, x_overlap = self.split_tiles(width, False)

        i_max, j_max = len(y_idx), len(x_idx)
        num_tiles = i_max * j_max

        x_tiles = []
        for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)):
            for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)):
                tile = x[..., i_pos : i_pos + i_len, j_pos : j_pos + j_len]
                x_tiles.append(tile)

        with ClsTokenAggregator(self) as agg:
            local_tile_indices = self._local_tile_indices(num_tiles, sp_rank, sp_size)
            stack_tiling = self.stack_tiling and not (
                self.training and getattr(self.encoder, "mask_enabled", False)
            )
            encoded_tasks = self._run_tile_tasks(
                x_tiles, local_tile_indices, self.encode, stack_tiling, agg
            )

            if sp_size > 1:
                dist.barrier(group=get_parallel_state()["sp_process_group"])
                all_encoded = self._all_gather_tiled_results(encoded_tasks, num_tiles)
                if agg.cls_tokens:
                    agg.cls_tokens = self._all_gather_tiled_results(agg.cls_tokens, num_tiles)
            else:
                all_encoded = encoded_tasks

        rows = [[None for _ in range(j_max)] for _ in range(i_max)]
        for idx, encoded in enumerate(all_encoded):
            i, j = idx // j_max, idx % j_max
            rows[i][j] = encoded.to(x.device)

        latent_y_overlap = [
            tile_overlap // self.vae_ratio for tile_overlap in y_overlap
        ]
        latent_x_overlap = [
            tile_overlap // self.vae_ratio for tile_overlap in x_overlap
        ]

        result_rows = []
        for i, row in enumerate(rows):
            result_row = []
            for j, tile in enumerate(row):
                if i > 0:
                    tile = self.blend(rows[i - 1][j], tile, latent_y_overlap[i - 1], dim=-2)
                if j > 0:
                    tile = self.blend(row[j - 1], tile, latent_x_overlap[j - 1], dim=-1)
                if i < len(rows) - 1:
                    tile = tile[..., : -latent_y_overlap[i], :]
                if j < len(row) - 1:
                    tile = tile[..., :, : -latent_x_overlap[j]]
                result_row.append(tile)
            result_rows.append(torch.cat(result_row, dim=-1))
        z = torch.cat(result_rows, dim=-2)

        return z

    def tiled_decode(self, z):
        if self.parallel_tiling:  # Fast online decoding for large videos
            state = get_parallel_state()
            sp_rank = state["sp_rank"]
            sp_size = state["sp_size"]
        else:
            sp_rank, sp_size = 0, 1

        height, width = (
            z.shape[-2] * self.vae_ratio,
            z.shape[-1] * self.vae_ratio,
        )
        y_idx, y_len, y_overlap = self.split_tiles(height, True)
        x_idx, x_len, x_overlap = self.split_tiles(width, True)

        i_max, j_max = len(y_idx), len(x_idx)
        num_tiles = i_max * j_max

        z_tiles = []
        for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)):
            i_pos, i_len = (
                i_pos // self.vae_ratio,
                i_len // self.vae_ratio,
            )
            for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)):
                j_pos, j_len = (j_pos // self.vae_ratio, j_len // self.vae_ratio)
                tile = z[..., i_pos : i_pos + i_len, j_pos : j_pos + j_len]
                z_tiles.append(tile)

        local_tile_indices = self._local_tile_indices(num_tiles, sp_rank, sp_size)
        stack_tiling = self.stack_tiling and not (
            self.training and getattr(self.decoder, "mask_enabled", False)
        )
        decoded_tasks = self._run_tile_tasks(
            z_tiles, local_tile_indices, self.decode, stack_tiling
        )

        if sp_size > 1:
            dist.barrier(group=get_parallel_state()["sp_process_group"])
            all_decoded = self._all_gather_tiled_results(decoded_tasks, num_tiles)
        else:
            all_decoded = decoded_tasks


        rows = [[None for _ in range(j_max)] for _ in range(i_max)]
        for idx, decoded in enumerate(all_decoded):
            i, j = idx // j_max, idx % j_max
            rows[i][j] = decoded.to(z.device)

        result_rows = []
        for i, row in enumerate(rows):
            result_row = []
            for j, tile in enumerate(row):
                if i > 0:
                    tile = self.blend(rows[i - 1][j], tile, y_overlap[i - 1], dim=-2)
                if j > 0:
                    tile = self.blend(row[j - 1], tile, x_overlap[j - 1], dim=-1)
                if i < len(rows) - 1:
                    tile = tile[..., : -y_overlap[i], :]
                if j < len(row) - 1:
                    tile = tile[..., :, : -x_overlap[j]]
                result_row.append(tile)
            result_rows.append(torch.cat(result_row, dim=-1))
        dec = torch.cat(result_rows, dim=-2)
        return dec

    def _adaptive_encode(self, x):
        if self.encoder_tiling:
            return self.tiled_encode(x)
        else:
            return self.encode(x)

    def _adaptive_decode(self, z):
        if self.decoder_tiling:
            return self.tiled_decode(z)
        else:
            return self.decode(z)

    def trim_code(self, z, target_codes):
        if target_codes < z.shape[2]:
            if self.causal_encoder:
                z = z[:, :, -target_codes:, :, :]
            else:
                start_frame = (z.shape[2] - target_codes) // 2
                z = z[:, :, start_frame : start_frame + target_codes, :, :]
        return z

    def trim_output(self, dec, target_frames):
        if target_frames < dec.shape[2]:
            if self.causal_encoder:  # This is defined by encoder, not decoder
                dec = dec[:, :, -target_frames:, :, :]
            else:
                start_frame = (dec.shape[2] - target_frames) // 2
                dec = dec[:, :, start_frame : start_frame + target_frames, :, :]
        return dec

    def encode_temporal(self, x):
        offset_frame = 1 if self.isolated_first_frame and self.frame_pre_padding == 0 else 0

        if x.shape[2] % self.clip_length != offset_frame:
            pad_size = (offset_frame - x.shape[2]) % self.clip_length
            pad_frames = x[:, :, -1:].repeat(1, 1, pad_size, 1, 1)
            x = torch.cat([x, pad_frames], dim=2)

        num_chunks = (x.shape[2] - offset_frame) // self.clip_length

        z_list = []
        for i in range(num_chunks):
            start_idx = i * self.clip_length + offset_frame
            end_idx = (i + 1) * self.clip_length + offset_frame
            clip_x = x[:, :, start_idx:end_idx, :, :]

            if self.isolated_key_frame:
                key_frame = clip_x[:, :, :1, :, :]
                z_key = self._adaptive_encode(key_frame)

                if clip_x.shape[2] > 1:
                    video_frames = clip_x[:, :, 1:, :, :]
                    z_video = self._adaptive_encode(video_frames)
                    z = torch.cat([z_key, z_video], dim=2)
                else:
                    z = z_key
            else:
                z = self._adaptive_encode(clip_x)

            z_list.append(z)

        z = torch.cat(z_list, dim=2)
        if self.token_drop > 0:
            z = z[:, :, : -self.token_drop]

        if self.isolated_first_frame:
            input_first_frame = x[:, :, :1, :, :]
            z_first_frame = self._adaptive_encode(input_first_frame)

            if self.frame_pre_padding == 0:
                z = torch.cat([z_first_frame, z], dim=2)
            else:
                z = torch.cat([z_first_frame, z[:, :, 1:, :, :]], dim=2)

        if self.isolated_last_frame:
            frame_num = x.shape[2]
            last_frame_idx = frame_num - self.frame_drop + offset_frame
            input_last_frame = x[:, :, last_frame_idx : last_frame_idx + 1, :, :]
            z_last_frame = self._adaptive_encode(input_last_frame)
            z = torch.cat([z, z_last_frame], dim=2)

        return z

    def _decode_temporal_pad_frames(self, z, pad_tokens):
        if pad_tokens <= 0:
            return 0
        intra_tail = self.clip_length % self.vae_ratio_t
        if intra_tail == 0:
            return int(pad_tokens) * int(self.vae_ratio_t)

        z_len_before_pad = z.shape[2] - pad_tokens
        return sum(
            (
                intra_tail
                if (z_len_before_pad + k) % self.tokens_chunk_size == 0
                else self.vae_ratio_t
            )
            for k in range(pad_tokens)
        )

    def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens):
        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
        split_count = int(self.token_drop > 0) + 1
        total_frames = 0
        final_overlap_frames = 0

        if z_head is not None:
            total_frames += 1

        for i in range(num_chunks):
            t_start_idx = i * self.tokens_chunk_size
            t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
            clip_token_len = max(0, min(t_end_idx, z.shape[2]) - min(t_start_idx, z.shape[2]))
            if i == 0 and z_head is not None:
                clip_token_len += z_head.shape[2]
            if i == num_chunks - 1 and z_tail is not None:
                clip_token_len += z_tail.shape[2]

            clip_frame_len = clip_token_len * self.vae_ratio_t
            if i == 0 and z_head is not None:
                clip_frame_len = max(0, clip_frame_len - self.vae_ratio_t)
            if i == num_chunks - 1 and z_tail is not None:
                clip_frame_len = max(0, clip_frame_len - self.vae_ratio_t)

            for j in range(split_count):
                f_start_idx = j * chunk_dec
                f_end_idx = min(f_start_idx + chunk_dec, clip_frame_len)
                chunk_frames = max(0, f_end_idx - f_start_idx - self.frame_pre_padding)
                if j == 0:
                    total_frames += chunk_frames
                else:
                    final_overlap_frames = chunk_frames

        total_frames += final_overlap_frames
        if z_tail is not None:
            total_frames += 1

        pad_frames = self._decode_temporal_pad_frames(z, pad_tokens)
        return int(total_frames), int(pad_frames), int(total_frames - pad_frames)

    def _decode_temporal_streaming(self, z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype):
        total_frames, pad_frames, output_frames = self._decode_temporal_output_frame_plan(
            z, z_head, z_tail, num_chunks, pad_tokens
        )
        if output_frames <= 0:
            raise ValueError(
                f"decode_temporal streaming planned non-positive output_frames={output_frames} "
                f"total_frames={total_frames} pad_frames={pad_frames}"
            )


        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
        split_count = int(self.token_drop > 0) + 1
        dec = None
        dec_overlap = None
        write_pos = 0
        logical_frames = 0
        dropped_frames = 0
        decoded_count = 0

        def write_part(part):
            nonlocal dec, write_pos, logical_frames, dropped_frames
            part_frames = int(part.shape[2])
            if part_frames <= 0:
                return
            logical_frames += part_frames
            if dec is None:
                out_shape = list(part.shape)
                out_shape[2] = output_frames
                dec = torch.empty(out_shape, dtype=part.dtype, device=part.device)

            remaining = int(dec.shape[2]) - write_pos
            copy_frames = min(part_frames, max(0, remaining))
            if copy_frames > 0:
                dec[:, :, write_pos : write_pos + copy_frames, :, :].copy_(
                    part[:, :, :copy_frames, :, :]
                )
                write_pos += copy_frames
            dropped_frames += part_frames - copy_frames

        for i in range(num_chunks):
            t_start_idx = i * self.tokens_chunk_size
            t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
            clip_z = z[:, :, t_start_idx:t_end_idx, :, :]

            if i == 0 and z_head is not None:
                clip_z = torch.cat([z_head, clip_z], dim=2)

            if i == num_chunks - 1 and z_tail is not None:
                clip_z = torch.cat([clip_z, z_tail], dim=2)

            clip_dec = self._adaptive_decode(clip_z)
            decoded_count += 1
            if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype:
                clip_dec = clip_dec.to(temporal_cat_dtype)
            if clip_dec.device != z.device:
                clip_dec = clip_dec.to(z.device)


            dec_tail = None
            if i == 0 and z_head is not None:
                write_part(clip_dec[:, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :])
                clip_dec = clip_dec[:, :, self.vae_ratio_t :, :, :]

            if i == num_chunks - 1 and z_tail is not None:
                dec_tail = clip_dec[:, :, -1:, :, :]
                clip_dec = clip_dec[:, :, : -self.vae_ratio_t, :, :]

            for j in range(split_count):
                f_start_idx = j * chunk_dec
                f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2])
                clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :]
                clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :]

                if j == 0:
                    if dec_overlap is not None:
                        clip_dec_chunk = self.blend(
                            dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3
                        )
                        dec_overlap = None
                    write_part(clip_dec_chunk)
                else:
                    # Break the view's reference to the full decoded clip so earlier
                    # temporal chunks can be released before the final output exists.
                    dec_overlap = clip_dec_chunk.contiguous()

            if i == num_chunks - 1:
                if dec_overlap is not None:
                    write_part(dec_overlap)
                    dec_overlap = None
                if dec_tail is not None:
                    write_part(dec_tail)

            del clip_dec, clip_z

        if dec is None:
            raise RuntimeError("decode_temporal streaming produced no output tensor")
        if logical_frames != total_frames or dropped_frames != pad_frames or write_pos != output_frames:
            raise RuntimeError(
                "decode_temporal streaming frame plan mismatch: "
                f"logical_frames={logical_frames} total_frames={total_frames} "
                f"dropped_frames={dropped_frames} pad_frames={pad_frames} "
                f"write_pos={write_pos} output_frames={output_frames}"
            )

        return dec

    def decode_temporal(self, z):
        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t

        isolated_token_num = 0
        if self.isolated_first_frame and self.frame_pre_padding == 0:
            isolated_token_num = isolated_token_num + 1
        if self.isolated_last_frame:
            isolated_token_num = isolated_token_num + 1

        pseudo_total_tokens = z.shape[2] - isolated_token_num + self.token_drop

        pad_tokens = 0
        remainder = pseudo_total_tokens % self.tokens_chunk_size
        if remainder != 0:
            if self.training:
                raise ValueError(f"Temporal token length {z.shape[2]} is wrong!")
            else:
                pad_tokens = self.tokens_chunk_size - remainder
                pseudo_total_tokens = pseudo_total_tokens + pad_tokens

        pseudo_num_chunks = pseudo_total_tokens // self.tokens_chunk_size
        num_chunks = pseudo_num_chunks - int(self.token_drop > 0)

        z_head = None
        if self.isolated_first_frame and self.frame_pre_padding == 0:
            z_head = z[:, :, :1, :, :]
            z = z[:, :, 1:, :, :]

        z_tail = None
        if self.isolated_last_frame:
            z_tail = z[:, :, -1:, :, :]
            z = z[:, :, :-1, :, :]

        if pad_tokens > 0:
            pad_z = z[:, :, -1:, :, :].repeat(1, 1, pad_tokens, 1, 1)
            z = torch.cat([z, pad_z], dim=2)

        temporal_cat_dtype = _resolve_temporal_cat_dtype()
        if not self.training and _resolve_temporal_stream_cat():
            return self._decode_temporal_streaming(
                z, z_head, z_tail, num_chunks, pad_tokens, temporal_cat_dtype
            )

        decoded_tasks = []
        for i in range(num_chunks):
            t_start_idx = i * self.tokens_chunk_size
            t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
            clip_z = z[:, :, t_start_idx:t_end_idx, :, :]

            if i == 0 and z_head is not None:
                clip_z = torch.cat([z_head, clip_z], dim=2)

            if i == num_chunks - 1 and z_tail is not None:
                clip_z = torch.cat([clip_z, z_tail], dim=2)

            clip_dec = self._adaptive_decode(clip_z)
            if temporal_cat_dtype is not None and clip_dec.dtype != temporal_cat_dtype:
                clip_dec = clip_dec.to(temporal_cat_dtype)

            decoded_tasks.append((i, clip_dec))

        clip_dec_list = [clip_dec.to(z.device) for _, clip_dec in decoded_tasks]

        dec_list = []
        dec_overlap = None

        dec_head = None
        if z_head is not None:
            dec_head = clip_dec_list[0][:, :, self.vae_ratio_t - 1 : self.vae_ratio_t, :, :]
            clip_dec_list[0] = clip_dec_list[0][:, :, self.vae_ratio_t :, :, :]

        dec_tail = None
        if z_tail is not None:
            dec_tail = clip_dec_list[-1][:, :, -1:, :, :]
            clip_dec_list[-1] = clip_dec_list[-1][:, :, : -self.vae_ratio_t, :, :]

        if dec_head is not None:
            dec_list.append(dec_head)

        for i in range(num_chunks):
            for j in range(int(self.token_drop > 0) + 1):
                clip_dec = clip_dec_list[i]

                f_start_idx = j * chunk_dec
                f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2])
                clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :]
                clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding :, :, :]

                if j == 0:
                    if dec_overlap is not None:
                        clip_dec_chunk = self.blend(
                            dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3
                        )
                    dec_list.append(clip_dec_chunk)
                else:
                    dec_overlap = clip_dec_chunk

        if dec_overlap is not None:
            dec_list.append(dec_overlap)

        if dec_tail is not None:
            dec_list.append(dec_tail)


        dec = torch.cat(dec_list, dim=2)

        pad_frames = self._decode_temporal_pad_frames(z, pad_tokens)
        if pad_frames > 0:
            dec = dec[:, :, :-pad_frames, :, :]

        return dec

    def decode_base(self, z, frame_num=None, process_image=False):
        if process_image or not self.use_3d_conv:
            if not self.use_3d_conv and z.ndim == 5:
                z = z.squeeze(2)

            recon = self._adaptive_decode(z)
        else:
            recon = self.decode_temporal(z)

        if self.use_3d_conv:
            if frame_num is not None:
                target_frames = frame_num
            else:
                target_frames = recon.shape[2]

            recon = self.trim_output(recon, target_frames)
            if process_image:
                recon = recon.squeeze(2)

        return recon

    #########################################################
    # freeze_scope is retained from the training codebase: in this
    # inference-only bundle (self.training is always False) it simply
    # provides the no_grad() context used by encode()/decode().
    #########################################################


    def freeze_scope(self, module_name):
        if not self.training:
            return torch.no_grad()

        if_freeze = module_name in self.fix_modules
        if if_freeze:
            return torch.no_grad()
        else:
            return nullcontext()








    #########################################################
    # following methods are for inference
    #########################################################

    @torch.no_grad()
    def encode_images(
        self,
        images: Union[List[np.ndarray], List[torch.Tensor]],
        transform_input: bool = False,
        use_fp16_latent: bool = False,
        verbose: bool = False,
    ) -> List[torch.Tensor]:
        """encode images into latents

        Args:
            images (Union[List[np.ndarray], List[torch.Tensor]]):
                List of images, single input will be wrapped in a list.
                If input is a list of np.ndarray, it should be in shape B * (H, W, 3), dtype uint8.
                If input is a list of torch.Tensor, it should be in shape B * (3, H, W), dtype float32.
            transform_input (bool, optional):
                Whether to transform input using ImageNet std/mean. Defaults to False.
                If input is a list of np.ndarray, it will always be set to True.
            use_fp16_latent (bool, optional):
                Whether to use fp16 latent. Defaults to False.
            verbose (bool, optional):
                Whether to print debug information. Defaults to False.

        Returns:
            List[torch.Tensor]:
                List of image latents.
                If self.use_3d_conv is True, it should be in shape B * (D, 1, H', W').
                Otherwise, it should be in shape B * (D, H', W').
        """

        images = self.processor._ensure_list(images)

        if isinstance(images[0], Image.Image):
            images = [np.array(image) for image in images]

        if isinstance(images[0], np.ndarray):
            device = next(self.parameters()).device
            images = self.processor.convert_numpy_to_tensor(images, device)
            images = torch.split(images, 1, dim=0)
            transform_input = True

        if transform_input:
            images = [
                image.unsqueeze(0) if image.ndim == 3 else image for image in images
            ]
            images = [self.processor.transform_tensor(image) for image in images]

        prepared = []
        for image_tensor in images:
            if image_tensor.ndim == 3:
                image_tensor = image_tensor.unsqueeze(0)
            _, _, h, w = image_tensor.shape
            new_h, new_w = self.processor._align_to_total_patch_size(h, w)
            image_tensor = self.processor._crop_to_align(image_tensor, new_h, new_w)
            prepared.append(image_tensor)

        if len(prepared) > 1 and len(set(t.shape for t in prepared)) == 1:
            stacked = torch.cat(prepared, dim=0)
            if verbose:
                logger.info(f"batch encode input shape {tuple(stacked.shape)}")
            all_latents = self.encode_base(stacked, True)
            image_latents = [all_latents[i].contiguous() for i in range(all_latents.shape[0])]
        else:
            image_latents = []
            for image_tensor in prepared:
                if verbose:
                    logger.info(f"input shape {tuple(image_tensor.shape)}")
                image_latent = self.encode_base(image_tensor, True)
                image_latents.append(image_latent.squeeze(0).contiguous())

        if use_fp16_latent:
            image_latents = [lat.to(torch.float16) for lat in image_latents]

        if verbose:
            for lat in image_latents:
                logger.info(f"image latent shape {tuple(lat.shape)}")

        return image_latents

    @torch.no_grad()
    def encode_videos(
        self,
        videos: Union[List[np.ndarray], List[torch.Tensor]],
        transform_input: bool = False,
        use_fp16_latent: bool = False,
        verbose: bool = False,
        encode_prefix: bool = False,
    ) -> List[torch.Tensor]:
        """encode videos into latents

        Args:
            videos (Union[List[np.ndarray], List[torch.Tensor]]):
                List of videos, single input will be wrapped in a list.
                If input is a list of np.ndarray, it should be in shape B * (T, H, W, 3), dtype uint8.
                If input is a list of torch.Tensor, it should be in shape B * (3, T, H, W), dtype float32.
            transform_input (bool, optional):
                Whether to transform input using ImageNet std/mean. Defaults to False.
                If input is a list of np.ndarray, it will always be set to True.
            use_fp16_latent (bool, optional):
                Whether to use fp16 latent. Defaults to False.
            verbose (bool, optional):
                Whether to print debug information. Defaults to False.
            encode_prefix (bool, optional):
                Continuation (prefix) mode: prepend normalized
                black frames to token alignment, append black frames to chunk
                alignment, encode with token_drop disabled, then discard only
                the trailing padding tokens. Returns both latents and leading
                pad-frame counts. Defaults to False.

        Returns:
            List[torch.Tensor]:
                List of video latents, shape B * (D, T', H', W').
                With encode_prefix=True, returns
                (List[torch.Tensor], List[int]).
        """

        videos = self.processor._ensure_list(videos)

        if isinstance(videos[0], np.ndarray):
            device = next(self.parameters()).device
            videos = [self.processor.convert_numpy_to_tensor(video, device) for video in videos]
            transform_input = True

        if transform_input:
            videos = [self.processor.transform_tensor(video) for video in videos]
            videos = [video.transpose(0, 1) for video in videos]

        if encode_prefix:
            if self.isolated_last_frame:
                raise ValueError(
                    "encode_prefix does not support isolated_last_frame"
                )

            video_latents = []
            prefix_pad_frames = []
            for video in videos:
                if video.ndim == 4:
                    video = video.unsqueeze(0)
                _, _, _, h, w = video.shape
                new_h, new_w = self.processor._align_to_total_patch_size(h, w)
                video = self.processor._crop_to_align(
                    video, new_h, new_w, is_video=True
                )

                model_alignment = (
                    self.token_drop,
                    self.frame_drop,
                    self.token_overlap,
                    self.frame_overlap,
                )
                processor_alignment = (
                    self.processor.token_overlap,
                    self.processor.frame_overlap,
                )
                self.token_drop = 0
                self.frame_drop = 0
                self.token_overlap = 0
                self.frame_overlap = 0
                self.processor.token_overlap = 0
                self.processor.frame_overlap = 0
                try:
                    orig_frames = video.shape[2]
                    leading, trailing, drop_tokens = (
                        self.processor.align_video_length_2pass(orig_frames)
                    )
                    _, _, _, cropped_h, cropped_w = video.shape
                    if leading > 0:
                        black = self.processor.transform(
                            video.new_zeros(leading, 3, cropped_h, cropped_w)
                        )
                        black = black.unsqueeze(0).permute(0, 2, 1, 3, 4)
                        video = torch.cat([black, video], dim=2)
                    if trailing > 0:
                        black = self.processor.transform(
                            video.new_zeros(trailing, 3, cropped_h, cropped_w)
                        )
                        black = black.unsqueeze(0).permute(0, 2, 1, 3, 4)
                        video = torch.cat([video, black], dim=2)

                    if verbose:
                        logger.info(
                            f"[encode_prefix] {orig_frames} frames -> "
                            f"pad leading={leading}, trailing={trailing} -> "
                            f"{video.shape[2]} frames"
                        )

                    video_latent = self.encode_base(video, False)
                    if drop_tokens > 0:
                        video_latent = video_latent[:, :, :-drop_tokens, :, :]
                    prefix_pad_frames.append(leading)
                finally:
                    (
                        self.token_drop,
                        self.frame_drop,
                        self.token_overlap,
                        self.frame_overlap,
                    ) = model_alignment
                    (
                        self.processor.token_overlap,
                        self.processor.frame_overlap,
                    ) = processor_alignment

                video_latents.append(video_latent.squeeze(0).contiguous())

            if use_fp16_latent:
                video_latents = [lat.to(torch.float16) for lat in video_latents]
            if verbose:
                for latent in video_latents:
                    logger.info(f"video latent shape {tuple(latent.shape)}")
            return video_latents, prefix_pad_frames

        prepared = []
        for video in videos:
            if video.ndim == 4:
                video = video.unsqueeze(0)
            used_frame_length = self.processor.get_suitable_video_length(video.shape[2], verbose)
            _, _, _, h, w = video.shape
            new_h, new_w = self.processor._align_to_total_patch_size(h, w)
            video = video[:, :, :used_frame_length, :, :]
            video = self.processor._crop_to_align(video, new_h, new_w, is_video=True)
            prepared.append(video)

        if len(prepared) > 1 and len(set(t.shape for t in prepared)) == 1:
            stacked = torch.cat(prepared, dim=0)
            if verbose:
                logger.info(f"batch encode input shape {tuple(stacked.shape)}")
            all_latents = self.encode_base(stacked, False)
            video_latents = [all_latents[i].contiguous() for i in range(all_latents.shape[0])]
        else:
            video_latents = []
            for video in prepared:
                if verbose:
                    logger.info(f"input shape {tuple(video.shape)}")
                video_latent = self.encode_base(video, False)
                video_latents.append(video_latent.squeeze(0).contiguous())

        if use_fp16_latent:
            video_latents = [lat.to(torch.float16) for lat in video_latents]

        if verbose:
            for lat in video_latents:
                logger.info(f"video latent shape {tuple(lat.shape)}")

        return video_latents




# ============================================================================
# Legacy CNN VAE
# ============================================================================


class AutoencoderKLLegacy(AutoencoderKL):
    r"""
    A VAE model (legacy CNN-based) for encoding pixels into latents and decoding latent representations into pixels.
    """

    @register_to_config
    def __init__(
        self,
        in_channels=3,
        out_ch=3,
        ch=128,
        embed_dim=16,
        z_channels=16,
        use_3d_conv=False,
        # cnn vae
        zq_ch_encoder=None,
        zq_ch_decoder=None,
        num_res_blocks=2,
        num_res_blocks_decoder=None,
        ch_mult=[1, 2, 2, 4, 4, 8],
        space_down=[2, 2, 2, 2, 1, 1],
        space_up=[1, 2, 2, 2, 2, 1],
        time_down=None,
        time_up=None,
        padding_mode="zeros",
        padding_mode_t=None,
        use_t_isolated_gn=False,
        causal_encoder=True,
        causal_decoder=True,
        use_vit_decoder=False,
        vit_decoder_kwargs=None,
        # stats
        shift_factor=0.0,
        scaling_factor=1.0,
        # pixel normalization
        pixel_norm_type="imagenet",
        # others
        **kwargs,
    ):
        ModelMixin.__init__(self)  # NOTE: avoid wrong @register_to_config

        if not use_3d_conv or not use_vit_decoder:
            raise NotImplementedError(
                "this release only supports use_3d_conv=True with use_vit_decoder=True"
            )

        self.transform = get_normalize_transform(pixel_norm_type)
        self.transform_rev = get_denormalize_transform(pixel_norm_type)

        self.use_3d_conv = use_3d_conv
        self.causal_encoder = causal_encoder
        self.causal_decoder = causal_decoder
        self.slidedec = self.causal_encoder and not self.causal_decoder

        # some registered parameters for simplicity
        self.vae_ratio = int(np.cumprod(space_down)[-1])
        self.vae_ratio_t = int(np.cumprod(time_down)[-1]) if time_down else 1
        self.config["vae_ratio"] = self.vae_ratio
        self.config["vae_ratio_t"] = self.vae_ratio_t

        # some registered parameters for inference and training
        self.setup_forward(**kwargs)
        self.setup_training(**kwargs)

        # init encoder
        encoder_config = {
            "double_z": True,
            "z_channels": z_channels,
            "zq_ch": zq_ch_encoder,
            "in_channels": in_channels,
            "ch": ch,
            "num_res_blocks": num_res_blocks,
            "ch_mult": ch_mult,
            "space_down": space_down,
            "time_down": time_down,
            "padding_mode": padding_mode,
            "padding_mode_t": padding_mode_t,
            "causal": causal_encoder,
            "use_t_isolated_gn": use_t_isolated_gn,
        }
        self.encoder = EncoderFCN3D(**encoder_config)

        # init pointwise quant/post_quant conv
        self.quant_conv = nn.Conv3d(z_channels * 2, 2 * embed_dim, 1)
        self.post_quant_conv = nn.Conv3d(embed_dim, z_channels, 1)

        self.use_vit_decoder = use_vit_decoder

        # init decoder
        vit_kwargs = {
            "patch_size": self.vae_ratio,
            "in_channels": z_channels,
            "out_channels": out_ch,
            **(vit_decoder_kwargs or {}),
        }
        vit_kwargs.setdefault("patch_size_t", self.vae_ratio_t)
        vit_kwargs.setdefault("t_causal", causal_decoder)
        self.decoder = ViT3DDecoder(**vit_kwargs)

        apply_spatial_parallel(self.encoder, self.encoder_parallel, self.chunk_dim)
        apply_spatial_parallel(self.decoder, self.decoder_parallel, self.chunk_dim)

        for module in set(self.fix_modules + self.frozen_modules):
            self._freeze_nested_module(module)

        self.gradient_checkpointing = False

    def encode(self, x):
        if self.encoder_parallel:
            x = self.perform_input_slice(x, self.vae_ratio)

        with self.freeze_scope("encoder"):
            h = self.encoder(x)

        with self.freeze_scope("quant_conv"):
            moments = self.quant_conv(h)

        if self.encoder_parallel:
            moments = self.perform_output_concat(moments)

        return moments

    def decode(self, z):
        if self.decoder_parallel and not self.use_vit_decoder:
            z = self.perform_input_slice(z)

        with self.freeze_scope("post_quant_conv"):
            z2 = self.post_quant_conv(z)

        with self.freeze_scope("decoder"):
            if self.use_vit_decoder:
                dec = self.decoder(z2)
            else:
                dec = self.decoder(z2, z)

        if self.decoder_parallel and not self.use_vit_decoder:
            dec = self.perform_output_concat(dec)
        return dec

    def encode_base(self, input, process_image=False):
        if self.use_3d_conv and input.ndim == 4:
            input = input.unsqueeze(2)

        if process_image or not self.use_3d_conv:
            moments = self._adaptive_encode(input)
        else:
            moments = self.encode_temporal(input)

        z = DiagonalGaussianDistribution(moments).sample()

        if process_image and self.use_3d_conv:
            z = self.trim_code(z, 1)

        return z

    #########################################################
    # training-related knobs kept only for checkpoint/config compatibility
    #########################################################

    def setup_training(self, **kwargs):
        self.fix_modules = kwargs.get("fix_modules", [])
        self.frozen_modules = kwargs.get("frozen_modules", [])