wehub-resource-sync
2296 lines
90 KiB
Python
2296 lines
90 KiB
Python
# Copyright 2025 The Lightricks team and The HuggingFace Team.
|
|
# All rights reserved.
|
|
#
|
|
# 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.
|
|
|
|
"""LTX-2 Video VAE with causal encoding/decoding and streaming cache support.
|
|
|
|
This module implements a 3D variational autoencoder for video compression with support for:
|
|
- Causal temporal convolution (future frames depend only on past frames)
|
|
- Streaming decode with persistent feature cache for memory-efficient long video processing
|
|
- Bidirectional encoding/decoding mode
|
|
|
|
Key classes:
|
|
- AutoencoderKLCausalLTX2Video: Main VAE model with cache management
|
|
- DecoderCacheManager: Manages decoder cache state for streaming inference
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Iterator
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
|
from diffusers.loaders import FromOriginalModelMixin
|
|
from diffusers.models.activations import get_activation
|
|
from diffusers.models.autoencoders.vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution
|
|
from diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings
|
|
from diffusers.models.modeling_outputs import AutoencoderKLOutput
|
|
from diffusers.models.modeling_utils import ModelMixin
|
|
from diffusers.utils.accelerate_utils import apply_forward_hook
|
|
|
|
# =============================================================================
|
|
# Utility Functions
|
|
# =============================================================================
|
|
|
|
|
|
def _shape_of(x) -> tuple | None:
|
|
"""Get shape of tensor if it is a torch.Tensor."""
|
|
if isinstance(x, torch.Tensor):
|
|
return tuple(x.shape)
|
|
return None
|
|
|
|
|
|
def _compute_conv_output_size(in_size: int, k: int, s: int, p: int, d: int) -> int:
|
|
"""Compute output size of a convolution operation."""
|
|
return max((in_size + 2 * p - d * (k - 1) - 1) // s + 1, 0)
|
|
|
|
|
|
# =============================================================================
|
|
# Cache Management Classes
|
|
# =============================================================================
|
|
|
|
|
|
@dataclass
|
|
class DecoderCacheState:
|
|
"""State container for decoder streaming cache.
|
|
|
|
Attributes:
|
|
feat_map: Per-layer feature cache for causal convolution
|
|
is_first_chunk: Whether this is the first chunk in the sequence
|
|
prev_latent_tail: Previous chunk's latent tail for context
|
|
cache_mode: Current cache mode (causal or not)
|
|
"""
|
|
|
|
feat_map: list = field(default_factory=list)
|
|
is_first_chunk: bool = True
|
|
prev_latent_tail: torch.Tensor | None = None
|
|
cache_mode: bool | None = None
|
|
|
|
|
|
class DecoderCacheManager:
|
|
"""Manages decoder cache state for streaming video decoding.
|
|
|
|
This class encapsulates all decoder cache logic, providing a clean interface
|
|
for cache operations used during streaming decode of long videos.
|
|
|
|
Example:
|
|
>>> manager = DecoderCacheManager()
|
|
>>> manager.clear() # Reset cache
|
|
>>> # In streaming loop:
|
|
>>> z_chunk = manager.prepend_context(z_chunk, prepend_frames=1)
|
|
>>> decoded = decoder(z_chunk, feat_cache=manager.feat_map, feat_idx=[0])
|
|
>>> decoded = manager.trim_output(decoded, prepend_frames=1, chunk_frames=z.shape[2])
|
|
>>> manager.update_tail(z_chunk, prepend_frames=1)
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._state = DecoderCacheState()
|
|
|
|
@property
|
|
def feat_map(self) -> list:
|
|
"""Get the feature cache map."""
|
|
return self._state.feat_map
|
|
|
|
@property
|
|
def is_first_chunk(self) -> bool:
|
|
"""Check if this is the first chunk."""
|
|
return self._state.is_first_chunk
|
|
|
|
@property
|
|
def prev_latent_tail(self) -> torch.Tensor | None:
|
|
"""Get previous latent tail."""
|
|
return self._state.prev_latent_tail
|
|
|
|
@property
|
|
def cache_mode(self) -> bool | None:
|
|
"""Get current cache mode."""
|
|
return self._state.cache_mode
|
|
|
|
def clear(self) -> None:
|
|
"""Clear all cache state."""
|
|
self._state = DecoderCacheState()
|
|
|
|
def validate_mode(self, causal: bool) -> None:
|
|
"""Validate and update cache mode.
|
|
|
|
Args:
|
|
causal: Desired causal mode
|
|
|
|
Raises:
|
|
ValueError: If mode conflicts with existing cache
|
|
"""
|
|
if self._state.cache_mode is None:
|
|
self._state.cache_mode = causal
|
|
elif self._state.cache_mode != causal:
|
|
# Mode mismatch - clear cache to avoid mixing states
|
|
self.clear()
|
|
self._state.cache_mode = causal
|
|
|
|
def prepend_context(
|
|
self, z: torch.Tensor, prepend_prev_latent_frames: int, temporal_compression_ratio: int
|
|
) -> torch.Tensor:
|
|
"""Prepend previous chunk's latent tail as left context.
|
|
|
|
Args:
|
|
z: Current latent chunk [B, C, T, H, W]
|
|
prepend_prev_latent_frames: Number of frames to prepend
|
|
temporal_compression_ratio: Ratio for converting latent to sample frames
|
|
|
|
Returns:
|
|
Latent tensor with context prepended
|
|
"""
|
|
if prepend_prev_latent_frames <= 0:
|
|
return z
|
|
|
|
prev_tail = self._state.prev_latent_tail
|
|
if prev_tail is None:
|
|
# First chunk: repeat first frame
|
|
left_ctx = z[:, :, :1, :, :].repeat(1, 1, prepend_prev_latent_frames, 1, 1)
|
|
else:
|
|
# Use previous chunk's tail
|
|
left_ctx = prev_tail.to(z.device)
|
|
if left_ctx.shape[2] > prepend_prev_latent_frames:
|
|
left_ctx = left_ctx[:, :, -prepend_prev_latent_frames:, :, :]
|
|
if left_ctx.shape[2] < prepend_prev_latent_frames:
|
|
fill = z[:, :, :1, :, :].repeat(1, 1, prepend_prev_latent_frames - left_ctx.shape[2], 1, 1)
|
|
left_ctx = torch.cat([fill, left_ctx], dim=2)
|
|
|
|
return torch.cat([left_ctx, z], dim=2)
|
|
|
|
def trim_output(
|
|
self,
|
|
decoded: torch.Tensor,
|
|
prepend_prev_latent_frames: int,
|
|
chunk_latent_frames: int,
|
|
temporal_compression_ratio: int,
|
|
) -> torch.Tensor:
|
|
"""Trim decoder output to remove context frames.
|
|
|
|
Args:
|
|
decoded: Full decoder output [B, C, T, H, W]
|
|
prepend_prev_latent_frames: Number of prepended latent frames
|
|
chunk_latent_frames: Number of latent frames in current chunk
|
|
temporal_compression_ratio: Ratio for converting latent to sample frames
|
|
|
|
Returns:
|
|
Trimmed output tensor
|
|
"""
|
|
drop_left_t = prepend_prev_latent_frames * temporal_compression_ratio
|
|
|
|
if self._state.is_first_chunk:
|
|
# First chunk: keep all frames from context start
|
|
keep_t = (chunk_latent_frames - 1) * temporal_compression_ratio + 1
|
|
self._state.is_first_chunk = False
|
|
else:
|
|
# Subsequent chunks: keep only new frames
|
|
keep_t = chunk_latent_frames * temporal_compression_ratio
|
|
|
|
return decoded[:, :, drop_left_t : drop_left_t + keep_t, :, :]
|
|
|
|
def update_tail(self, z: torch.Tensor, prepend_prev_latent_frames: int) -> None:
|
|
"""Update the previous latent tail for next chunk.
|
|
|
|
Args:
|
|
z: Current latent chunk before context prepending
|
|
prepend_prev_latent_frames: Number of frames to save as tail
|
|
"""
|
|
if prepend_prev_latent_frames > 0:
|
|
self._state.prev_latent_tail = z[:, :, -prepend_prev_latent_frames:, :, :].clone()
|
|
|
|
|
|
class EncoderCacheManager:
|
|
"""Manages encoder cache state for streaming video encoding."""
|
|
|
|
def __init__(self):
|
|
self._feat_map: list = []
|
|
|
|
@property
|
|
def feat_map(self) -> list:
|
|
"""Get the feature cache map."""
|
|
return self._feat_map
|
|
|
|
def clear(self) -> None:
|
|
"""Clear all cache state."""
|
|
self._feat_map = []
|
|
|
|
def check_pending_consumed(self) -> None:
|
|
"""Verify that all temporal grouping state has been fully consumed."""
|
|
for state in self._feat_map:
|
|
if isinstance(state, dict) and state.get("pending") is not None:
|
|
pending = state["pending"]
|
|
if isinstance(pending, torch.Tensor) and pending.shape[2] > 0:
|
|
raise RuntimeError(
|
|
"Encoder ended with non-empty temporal pending state. "
|
|
"Use chunk sizes aligned with temporal downsampling or process more frames."
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Normalization Layers
|
|
# =============================================================================
|
|
|
|
|
|
class PerChannelRMSNorm(nn.Module):
|
|
"""Per-pixel (per-location) RMS normalization layer.
|
|
|
|
For each element along the chosen dimension, this layer normalizes the tensor by the root-mean-square of its values
|
|
across that dimension:
|
|
|
|
y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps)
|
|
|
|
Args:
|
|
channel_dim: Dimension along which to compute the RMS (typically channels).
|
|
eps: Small constant added for numerical stability.
|
|
"""
|
|
|
|
def __init__(self, channel_dim: int = 1, eps: float = 1e-8) -> None:
|
|
super().__init__()
|
|
self.channel_dim = channel_dim
|
|
self.eps = eps
|
|
|
|
def forward(self, x: torch.Tensor, channel_dim: int | None = None) -> torch.Tensor:
|
|
"""Apply RMS normalization along the configured dimension."""
|
|
channel_dim = channel_dim or self.channel_dim
|
|
mean_sq = torch.mean(x**2, dim=self.channel_dim, keepdim=True)
|
|
rms = torch.sqrt(mean_sq + self.eps)
|
|
return x / rms
|
|
|
|
|
|
# =============================================================================
|
|
# Causal Convolution Layer
|
|
# =============================================================================
|
|
|
|
|
|
class LTX2VideoCausalConv3d(nn.Module):
|
|
"""Causal 3D convolution for video processing.
|
|
|
|
Like LTXCausalConv3d, but whether causal inference is performed can be
|
|
specified at runtime via the `causal` parameter.
|
|
|
|
Args:
|
|
in_channels: Number of input channels
|
|
out_channels: Number of output channels
|
|
kernel_size: Temporal, height, width kernel size (int or tuple)
|
|
stride: Stride for convolution (int or tuple)
|
|
dilation: Dilation for convolution (int or tuple)
|
|
groups: Number of groups for grouped convolution
|
|
spatial_padding_mode: Padding mode for spatial dimensions
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int | tuple[int, int, int] = 3,
|
|
stride: int | tuple[int, int, int] = 1,
|
|
dilation: int | tuple[int, int, int] = 1,
|
|
groups: int = 1,
|
|
spatial_padding_mode: str = "zeros",
|
|
):
|
|
super().__init__()
|
|
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size, kernel_size)
|
|
|
|
dilation = dilation if isinstance(dilation, tuple) else (dilation, 1, 1)
|
|
stride = stride if isinstance(stride, tuple) else (stride, stride, stride)
|
|
height_pad = self.kernel_size[1] // 2
|
|
width_pad = self.kernel_size[2] // 2
|
|
padding = (0, height_pad, width_pad)
|
|
|
|
self.conv = nn.Conv3d(
|
|
in_channels,
|
|
out_channels,
|
|
self.kernel_size,
|
|
stride=stride,
|
|
dilation=dilation,
|
|
groups=groups,
|
|
padding=padding,
|
|
padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
causal: bool = True,
|
|
feat_cache: list | None = None,
|
|
feat_idx: list | None = None,
|
|
) -> torch.Tensor:
|
|
"""Forward pass with optional causal convolution and feature caching.
|
|
|
|
Args:
|
|
hidden_states: Input tensor [B, C, T, H, W]
|
|
causal: Whether to use causal convolution
|
|
feat_cache: Cache for temporal feature reuse
|
|
feat_idx: Current index into feat_cache (mutated in-place)
|
|
|
|
Returns:
|
|
Output tensor after convolution
|
|
"""
|
|
time_kernel_size = self.kernel_size[0]
|
|
input_shape = _shape_of(hidden_states)
|
|
|
|
# Handle empty temporal chunks
|
|
if hidden_states.shape[2] == 0:
|
|
return self._handle_empty_chunk(hidden_states, causal, feat_cache, feat_idx, time_kernel_size)
|
|
|
|
if causal:
|
|
hidden_states = self._apply_causal_padding(
|
|
hidden_states, feat_cache, feat_idx, time_kernel_size, input_shape
|
|
)
|
|
else:
|
|
hidden_states = self._apply_bidirectional_padding(hidden_states, time_kernel_size)
|
|
|
|
hidden_states = self.conv(hidden_states)
|
|
return hidden_states
|
|
|
|
def _handle_empty_chunk(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
causal: bool,
|
|
feat_cache: list | None,
|
|
feat_idx: list | None,
|
|
time_kernel_size: int,
|
|
) -> torch.Tensor:
|
|
"""Handle empty temporal chunks by reserving cache slots."""
|
|
cache_len = max(time_kernel_size - 1, 0)
|
|
if causal and feat_cache is not None and feat_idx is not None and cache_len > 0:
|
|
idx = feat_idx[0]
|
|
while len(feat_cache) <= idx:
|
|
feat_cache.append(None)
|
|
feat_idx[0] += 1
|
|
|
|
batch = hidden_states.shape[0]
|
|
out_h = _compute_conv_output_size(
|
|
hidden_states.shape[3],
|
|
self.conv.kernel_size[1],
|
|
self.conv.stride[1],
|
|
self.conv.padding[1],
|
|
self.conv.dilation[1],
|
|
)
|
|
out_w = _compute_conv_output_size(
|
|
hidden_states.shape[4],
|
|
self.conv.kernel_size[2],
|
|
self.conv.stride[2],
|
|
self.conv.padding[2],
|
|
self.conv.dilation[2],
|
|
)
|
|
return hidden_states.new_empty(batch, self.conv.out_channels, 0, out_h, out_w)
|
|
|
|
def _apply_causal_padding(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
feat_cache: list | None,
|
|
feat_idx: list | None,
|
|
time_kernel_size: int,
|
|
input_shape: tuple | None,
|
|
) -> torch.Tensor:
|
|
"""Apply causal padding using cached history."""
|
|
cache_len = max(time_kernel_size - 1, 0)
|
|
if feat_cache is not None and feat_idx is not None and cache_len > 0:
|
|
idx = feat_idx[0]
|
|
while len(feat_cache) <= idx:
|
|
feat_cache.append(None)
|
|
|
|
prefix = feat_cache[idx]
|
|
current_cache = hidden_states[:, :, -cache_len:, :, :].clone()
|
|
|
|
# Handle short chunks by preserving history
|
|
if isinstance(prefix, torch.Tensor) and current_cache.shape[2] < cache_len:
|
|
needed = cache_len - current_cache.shape[2]
|
|
carry = prefix.to(hidden_states.device)
|
|
carry = carry[:, :, -needed:, :, :]
|
|
if carry.shape[2] < needed:
|
|
fill = carry[:, :, :1, :, :].repeat((1, 1, needed - carry.shape[2], 1, 1))
|
|
carry = torch.cat([fill, carry], dim=2)
|
|
current_cache = torch.cat([carry, current_cache], dim=2)
|
|
|
|
# Prepare prefix from cache or fallback
|
|
if prefix is not None:
|
|
prefix = prefix.to(hidden_states.device)
|
|
if prefix.shape[2] > cache_len:
|
|
prefix = prefix[:, :, -cache_len:, :, :]
|
|
if prefix.shape[2] < cache_len:
|
|
# Extend by repeating oldest cached frame
|
|
fallback = prefix[:, :, :1, :, :].repeat((1, 1, cache_len - prefix.shape[2], 1, 1))
|
|
prefix = torch.concatenate([fallback, prefix], dim=2)
|
|
else:
|
|
# No cache: use first frame as fallback
|
|
prefix = hidden_states[:, :, :1, :, :].repeat((1, 1, cache_len, 1, 1))
|
|
|
|
hidden_states = torch.concatenate([prefix, hidden_states], dim=2)
|
|
feat_cache[idx] = current_cache
|
|
feat_idx[0] += 1
|
|
else:
|
|
# No cache: use causal padding
|
|
pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, cache_len, 1, 1))
|
|
hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
|
|
|
|
return hidden_states
|
|
|
|
def _apply_bidirectional_padding(self, hidden_states: torch.Tensor, time_kernel_size: int) -> torch.Tensor:
|
|
"""Apply symmetric padding for bidirectional mode."""
|
|
pad_size = (time_kernel_size - 1) // 2
|
|
pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, pad_size, 1, 1))
|
|
pad_right = hidden_states[:, :, -1:, :, :].repeat((1, 1, pad_size, 1, 1))
|
|
return torch.concatenate([pad_left, hidden_states, pad_right], dim=2)
|
|
|
|
|
|
# =============================================================================
|
|
# ResNet Block
|
|
# =============================================================================
|
|
|
|
|
|
class LTX2VideoResnetBlock3d(nn.Module):
|
|
"""A 3D ResNet block used in the LTX 2.0 audiovisual model.
|
|
|
|
Args:
|
|
in_channels: Number of input channels.
|
|
out_channels: Number of output channels. If None, defaults to `in_channels`.
|
|
dropout: Dropout rate.
|
|
eps: Epsilon value for normalization layers.
|
|
elementwise_affine: Whether to enable elementwise affinity in normalization.
|
|
non_linearity: Activation function to use.
|
|
conv_shortcut: Whether to use a convolution shortcut.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int | None = None,
|
|
dropout: float = 0.0,
|
|
eps: float = 1e-6,
|
|
elementwise_affine: bool = False,
|
|
non_linearity: str = "swish",
|
|
inject_noise: bool = False,
|
|
timestep_conditioning: bool = False,
|
|
spatial_padding_mode: str = "zeros",
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
out_channels = out_channels or in_channels
|
|
|
|
self.nonlinearity = get_activation(non_linearity)
|
|
|
|
self.norm1 = PerChannelRMSNorm()
|
|
self.conv1 = LTX2VideoCausalConv3d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=3,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
self.norm2 = PerChannelRMSNorm()
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.conv2 = LTX2VideoCausalConv3d(
|
|
in_channels=out_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=3,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
self.norm3 = None
|
|
self.conv_shortcut = None
|
|
if in_channels != out_channels:
|
|
self.norm3 = nn.LayerNorm(in_channels, eps=eps, elementwise_affine=True, bias=True)
|
|
# LTX 2.0 uses a normal nn.Conv3d here rather than LTXVideoCausalConv3d
|
|
self.conv_shortcut = nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1)
|
|
|
|
self.per_channel_scale1 = None
|
|
self.per_channel_scale2 = None
|
|
if inject_noise:
|
|
self.per_channel_scale1 = nn.Parameter(torch.zeros(in_channels, 1, 1))
|
|
self.per_channel_scale2 = nn.Parameter(torch.zeros(in_channels, 1, 1))
|
|
|
|
self.scale_shift_table = None
|
|
if timestep_conditioning:
|
|
self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5)
|
|
|
|
def forward(
|
|
self,
|
|
inputs: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
generator: torch.Generator | None = None,
|
|
causal: bool = True,
|
|
feat_cache: list | None = None,
|
|
feat_idx: list | None = None,
|
|
) -> torch.Tensor:
|
|
"""Forward pass through ResNet block."""
|
|
hidden_states = inputs
|
|
|
|
hidden_states = self.norm1(hidden_states)
|
|
|
|
if self.scale_shift_table is not None:
|
|
temb = temb.unflatten(1, (4, -1)) + self.scale_shift_table[None, ..., None, None, None]
|
|
shift_1, scale_1, shift_2, scale_2 = temb.unbind(dim=1)
|
|
hidden_states = hidden_states * (1 + scale_1) + shift_1
|
|
|
|
hidden_states = self.nonlinearity(hidden_states)
|
|
hidden_states = self.conv1(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
if self.per_channel_scale1 is not None:
|
|
spatial_shape = hidden_states.shape[-2:]
|
|
spatial_noise = torch.randn(
|
|
spatial_shape, generator=generator, device=hidden_states.device, dtype=hidden_states.dtype
|
|
)[None]
|
|
hidden_states = hidden_states + (spatial_noise * self.per_channel_scale1)[None, :, None, ...]
|
|
|
|
hidden_states = self.norm2(hidden_states)
|
|
|
|
if self.scale_shift_table is not None:
|
|
hidden_states = hidden_states * (1 + scale_2) + shift_2
|
|
|
|
hidden_states = self.nonlinearity(hidden_states)
|
|
hidden_states = self.dropout(hidden_states)
|
|
hidden_states = self.conv2(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
if self.per_channel_scale2 is not None:
|
|
spatial_shape = hidden_states.shape[-2:]
|
|
spatial_noise = torch.randn(
|
|
spatial_shape, generator=generator, device=hidden_states.device, dtype=hidden_states.dtype
|
|
)[None]
|
|
hidden_states = hidden_states + (spatial_noise * self.per_channel_scale2)[None, :, None, ...]
|
|
|
|
if self.norm3 is not None:
|
|
inputs = self.norm3(inputs.movedim(1, -1)).movedim(-1, 1)
|
|
|
|
if self.conv_shortcut is not None:
|
|
inputs = self.conv_shortcut(inputs)
|
|
|
|
hidden_states = hidden_states + inputs
|
|
return hidden_states
|
|
|
|
|
|
# =============================================================================
|
|
# Downsampler and Upsampler
|
|
# =============================================================================
|
|
|
|
|
|
class LTXVideoDownsampler3d(nn.Module):
|
|
"""3D downsampling layer for spatiotemporal reduction.
|
|
|
|
Uses pixel unshuffle pattern for downsampling with causal temporal handling.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
stride: int | tuple[int, int, int] = 1,
|
|
spatial_padding_mode: str = "zeros",
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride)
|
|
self.group_size = (in_channels * stride[0] * stride[1] * stride[2]) // out_channels
|
|
|
|
out_channels = out_channels // (self.stride[0] * self.stride[1] * self.stride[2])
|
|
|
|
self.conv = LTX2VideoCausalConv3d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=3,
|
|
stride=1,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
causal: bool = True,
|
|
feat_cache: list | None = None,
|
|
feat_idx: list | None = None,
|
|
) -> torch.Tensor:
|
|
"""Forward pass with temporal grouping state management."""
|
|
_shape_of(hidden_states)
|
|
prefix_len = max(self.stride[0] - 1, 0)
|
|
|
|
if prefix_len > 0 and feat_cache is not None and feat_idx is not None:
|
|
hidden_states = self._manage_temporal_grouping(hidden_states, feat_cache, feat_idx, prefix_len)
|
|
else:
|
|
hidden_states = torch.cat([hidden_states[:, :, : self.stride[0] - 1], hidden_states], dim=2)
|
|
|
|
# Handle empty temporal dimension after grouping
|
|
if hidden_states.shape[2] == 0:
|
|
return self._handle_empty_output(hidden_states, feat_cache, feat_idx)
|
|
|
|
# Apply pixel unshuffle downsampling
|
|
residual = self._compute_residual(hidden_states)
|
|
|
|
hidden_states = self.conv(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
hidden_states = self._unshuffle_output(hidden_states)
|
|
hidden_states = hidden_states + residual
|
|
|
|
return hidden_states
|
|
|
|
def _manage_temporal_grouping(
|
|
self, hidden_states: torch.Tensor, feat_cache: list, feat_idx: list, prefix_len: int
|
|
) -> torch.Tensor:
|
|
"""Manage temporal grouping state for chunked processing."""
|
|
idx = feat_idx[0]
|
|
while len(feat_cache) <= idx:
|
|
feat_cache.append(None)
|
|
|
|
state = feat_cache[idx]
|
|
if state is None or not isinstance(state, dict):
|
|
state = {"prefixed": False, "pending": None}
|
|
|
|
sequence = hidden_states
|
|
if not state["prefixed"]:
|
|
prefix = hidden_states[:, :, :1, :, :].repeat((1, 1, prefix_len, 1, 1))
|
|
sequence = torch.cat([prefix, sequence], dim=2)
|
|
state["prefixed"] = True
|
|
|
|
if state["pending"] is not None:
|
|
pending = state["pending"].to(sequence.device)
|
|
sequence = torch.cat([pending, sequence], dim=2)
|
|
|
|
stride_t = self.stride[0]
|
|
usable_t = (sequence.shape[2] // stride_t) * stride_t
|
|
state["pending"] = sequence[:, :, usable_t:, :, :].clone() if usable_t < sequence.shape[2] else None
|
|
hidden_states = sequence[:, :, :usable_t, :, :]
|
|
feat_cache[idx] = state
|
|
feat_idx[0] += 1
|
|
|
|
return hidden_states
|
|
|
|
def _handle_empty_output(
|
|
self, hidden_states: torch.Tensor, feat_cache: list | None, feat_idx: list | None
|
|
) -> torch.Tensor:
|
|
"""Handle empty output when temporal dimension is 0."""
|
|
if feat_cache is not None and feat_idx is not None:
|
|
# Reserve cache slot for conv
|
|
idx = feat_idx[0]
|
|
while len(feat_cache) <= idx:
|
|
feat_cache.append(None)
|
|
feat_idx[0] += 1
|
|
|
|
out_channels = self.conv.out_channels * self.stride[0] * self.stride[1] * self.stride[2]
|
|
out_h = hidden_states.shape[3] // self.stride[1]
|
|
out_w = hidden_states.shape[4] // self.stride[2]
|
|
return hidden_states.new_empty(
|
|
hidden_states.shape[0],
|
|
out_channels,
|
|
0,
|
|
out_h,
|
|
out_w,
|
|
)
|
|
|
|
def _compute_residual(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
"""Compute residual connection for pixel unshuffle."""
|
|
residual = (
|
|
hidden_states.unflatten(4, (-1, self.stride[2]))
|
|
.unflatten(3, (-1, self.stride[1]))
|
|
.unflatten(2, (-1, self.stride[0]))
|
|
)
|
|
residual = residual.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(1, 4)
|
|
residual = residual.unflatten(1, (-1, self.group_size))
|
|
return residual.mean(dim=2)
|
|
|
|
def _unshuffle_output(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
"""Apply pixel unshuffle to output."""
|
|
hidden_states = (
|
|
hidden_states.unflatten(4, (-1, self.stride[2]))
|
|
.unflatten(3, (-1, self.stride[1]))
|
|
.unflatten(2, (-1, self.stride[0]))
|
|
)
|
|
hidden_states = hidden_states.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(1, 4)
|
|
return hidden_states
|
|
|
|
|
|
class LTXVideoUpsampler3d(nn.Module):
|
|
"""3D upsampling layer for spatiotemporal expansion.
|
|
|
|
Uses pixel shuffle pattern for upsampling with causal temporal handling.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
stride: int | tuple[int, int, int] = 1,
|
|
residual: bool = False,
|
|
upscale_factor: int = 1,
|
|
spatial_padding_mode: str = "zeros",
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride)
|
|
self.residual = residual
|
|
self.upscale_factor = upscale_factor
|
|
|
|
out_channels = (in_channels * stride[0] * stride[1] * stride[2]) // upscale_factor
|
|
|
|
self.conv = LTX2VideoCausalConv3d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=3,
|
|
stride=1,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
causal: bool = True,
|
|
feat_cache: list | None = None,
|
|
feat_idx: list | None = None,
|
|
) -> torch.Tensor:
|
|
"""Forward pass with trim state management."""
|
|
batch_size, num_channels, num_frames, height, width = hidden_states.shape
|
|
_shape_of(hidden_states)
|
|
trim_t = max(self.stride[0] - 1, 0)
|
|
|
|
# Determine trim amount based on cache state
|
|
if trim_t > 0 and feat_cache is not None and feat_idx is not None:
|
|
trim_start = self._manage_trim_state(feat_cache, feat_idx, trim_t)
|
|
else:
|
|
trim_start = trim_t
|
|
|
|
# Compute residual if enabled
|
|
if self.residual:
|
|
residual = self._compute_residual(hidden_states, trim_start)
|
|
|
|
hidden_states = self.conv(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
hidden_states = self._shuffle_output(hidden_states, trim_start)
|
|
|
|
if self.residual:
|
|
hidden_states = hidden_states + residual
|
|
|
|
return hidden_states
|
|
|
|
def _manage_trim_state(self, feat_cache: list, feat_idx: list, trim_t: int) -> int:
|
|
"""Manage trim state for removing temporal padding."""
|
|
idx = feat_idx[0]
|
|
while len(feat_cache) <= idx:
|
|
feat_cache.append(None)
|
|
state = feat_cache[idx]
|
|
if state is None or not isinstance(state, dict):
|
|
state = {"trim_applied": False}
|
|
trim_start = trim_t if not state["trim_applied"] else 0
|
|
state["trim_applied"] = True
|
|
feat_cache[idx] = state
|
|
feat_idx[0] += 1
|
|
return trim_start
|
|
|
|
def _compute_residual(self, hidden_states: torch.Tensor, trim_start: int) -> torch.Tensor:
|
|
"""Compute residual connection."""
|
|
batch_size, num_channels, num_frames, height, width = hidden_states.shape
|
|
residual = hidden_states.reshape(
|
|
batch_size, -1, self.stride[0], self.stride[1], self.stride[2], num_frames, height, width
|
|
)
|
|
residual = residual.permute(0, 1, 5, 2, 6, 3, 7, 4).flatten(6, 7).flatten(4, 5).flatten(2, 3)
|
|
repeats = (self.stride[0] * self.stride[1] * self.stride[2]) // self.upscale_factor
|
|
residual = residual.repeat(1, repeats, 1, 1, 1)
|
|
return residual[:, :, trim_start:]
|
|
|
|
def _shuffle_output(self, hidden_states: torch.Tensor, trim_start: int) -> torch.Tensor:
|
|
"""Apply pixel shuffle and trim to output."""
|
|
batch_size, num_channels, num_frames, height, width = hidden_states.shape
|
|
hidden_states = hidden_states.reshape(
|
|
batch_size, -1, self.stride[0], self.stride[1], self.stride[2], num_frames, height, width
|
|
)
|
|
hidden_states = hidden_states.permute(0, 1, 5, 2, 6, 3, 7, 4).flatten(6, 7).flatten(4, 5).flatten(2, 3)
|
|
return hidden_states[:, :, trim_start:]
|
|
|
|
|
|
# =============================================================================
|
|
# Encoder/Decoder Blocks
|
|
# =============================================================================
|
|
|
|
|
|
class LTX2VideoDownBlock3D(nn.Module):
|
|
"""Down block used in the LTXVideo model.
|
|
|
|
Args:
|
|
in_channels: Number of input channels.
|
|
out_channels: Number of output channels. If None, defaults to `in_channels`.
|
|
num_layers: Number of resnet layers.
|
|
dropout: Dropout rate.
|
|
resnet_eps: Epsilon value for normalization layers.
|
|
resnet_act_fn: Activation function to use.
|
|
spatio_temporal_scale: Whether to use downsampling layer.
|
|
downsample_type: Type of downsampling ("conv", "spatial", "temporal", "spatiotemporal")
|
|
"""
|
|
|
|
_supports_gradient_checkpointing = True
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int | None = None,
|
|
num_layers: int = 1,
|
|
dropout: float = 0.0,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_act_fn: str = "swish",
|
|
spatio_temporal_scale: bool = True,
|
|
downsample_type: str = "conv",
|
|
spatial_padding_mode: str = "zeros",
|
|
):
|
|
super().__init__()
|
|
|
|
out_channels = out_channels or in_channels
|
|
|
|
resnets = []
|
|
for _ in range(num_layers):
|
|
resnets.append(
|
|
LTX2VideoResnetBlock3d(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
dropout=dropout,
|
|
eps=resnet_eps,
|
|
non_linearity=resnet_act_fn,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
self.downsamplers = None
|
|
if spatio_temporal_scale:
|
|
self.downsamplers = nn.ModuleList()
|
|
|
|
if downsample_type == "conv":
|
|
self.downsamplers.append(
|
|
LTX2VideoCausalConv3d(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
kernel_size=3,
|
|
stride=(2, 2, 2),
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
)
|
|
elif downsample_type == "spatial":
|
|
self.downsamplers.append(
|
|
LTXVideoDownsampler3d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
stride=(1, 2, 2),
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
)
|
|
elif downsample_type == "temporal":
|
|
self.downsamplers.append(
|
|
LTXVideoDownsampler3d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
stride=(2, 1, 1),
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
)
|
|
elif downsample_type == "spatiotemporal":
|
|
self.downsamplers.append(
|
|
LTXVideoDownsampler3d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
stride=(2, 2, 2),
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
)
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
generator: torch.Generator | None = None,
|
|
causal: bool = True,
|
|
feat_cache: list | None = None,
|
|
feat_idx: list | None = None,
|
|
) -> torch.Tensor:
|
|
for i, resnet in enumerate(self.resnets):
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing and feat_cache is None:
|
|
hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, generator, causal)
|
|
else:
|
|
hidden_states = resnet(
|
|
hidden_states, temb, generator, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx
|
|
)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class LTX2VideoMidBlock3d(nn.Module):
|
|
"""Middle block used in the LTXVideo model.
|
|
|
|
Args:
|
|
in_channels: Number of input channels.
|
|
num_layers: Number of resnet layers.
|
|
dropout: Dropout rate.
|
|
resnet_eps: Epsilon value for normalization layers.
|
|
resnet_act_fn: Activation function to use.
|
|
"""
|
|
|
|
_supports_gradient_checkpointing = True
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
num_layers: int = 1,
|
|
dropout: float = 0.0,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_act_fn: str = "swish",
|
|
inject_noise: bool = False,
|
|
timestep_conditioning: bool = False,
|
|
spatial_padding_mode: str = "zeros",
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.time_embedder = None
|
|
if timestep_conditioning:
|
|
self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(in_channels * 4, 0)
|
|
|
|
resnets = []
|
|
for _ in range(num_layers):
|
|
resnets.append(
|
|
LTX2VideoResnetBlock3d(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
dropout=dropout,
|
|
eps=resnet_eps,
|
|
non_linearity=resnet_act_fn,
|
|
inject_noise=inject_noise,
|
|
timestep_conditioning=timestep_conditioning,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
generator: torch.Generator | None = None,
|
|
causal: bool = True,
|
|
feat_cache: list | None = None,
|
|
feat_idx: list | None = None,
|
|
) -> torch.Tensor:
|
|
if self.time_embedder is not None:
|
|
temb = self.time_embedder(
|
|
timestep=temb.flatten(),
|
|
resolution=None,
|
|
aspect_ratio=None,
|
|
batch_size=hidden_states.size(0),
|
|
hidden_dtype=hidden_states.dtype,
|
|
)
|
|
temb = temb.view(hidden_states.size(0), -1, 1, 1, 1)
|
|
|
|
for i, resnet in enumerate(self.resnets):
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing and feat_cache is None:
|
|
hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, generator, causal)
|
|
else:
|
|
hidden_states = resnet(
|
|
hidden_states, temb, generator, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx
|
|
)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class LTX2VideoUpBlock3d(nn.Module):
|
|
"""Up block used in the LTXVideo model.
|
|
|
|
Args:
|
|
in_channels: Number of input channels.
|
|
out_channels: Number of output channels. If None, defaults to `in_channels`.
|
|
num_layers: Number of resnet layers.
|
|
dropout: Dropout rate.
|
|
resnet_eps: Epsilon value for normalization layers.
|
|
resnet_act_fn: Activation function to use.
|
|
spatio_temporal_scale: Whether to use upsampling layer.
|
|
"""
|
|
|
|
_supports_gradient_checkpointing = True
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int | None = None,
|
|
num_layers: int = 1,
|
|
dropout: float = 0.0,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_act_fn: str = "swish",
|
|
spatio_temporal_scale: bool = True,
|
|
inject_noise: bool = False,
|
|
timestep_conditioning: bool = False,
|
|
upsample_residual: bool = False,
|
|
upscale_factor: int = 1,
|
|
spatial_padding_mode: str = "zeros",
|
|
upsample_stride: tuple[int, int, int] = (2, 2, 2),
|
|
):
|
|
super().__init__()
|
|
|
|
out_channels = out_channels or in_channels
|
|
|
|
self.time_embedder = None
|
|
if timestep_conditioning:
|
|
self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(in_channels * 4, 0)
|
|
|
|
self.conv_in = None
|
|
if in_channels != out_channels:
|
|
self.conv_in = LTX2VideoResnetBlock3d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
dropout=dropout,
|
|
eps=resnet_eps,
|
|
non_linearity=resnet_act_fn,
|
|
inject_noise=inject_noise,
|
|
timestep_conditioning=timestep_conditioning,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
self.upsamplers = None
|
|
if spatio_temporal_scale:
|
|
self.upsamplers = nn.ModuleList(
|
|
[
|
|
LTXVideoUpsampler3d(
|
|
out_channels * upscale_factor,
|
|
stride=upsample_stride,
|
|
residual=upsample_residual,
|
|
upscale_factor=upscale_factor,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
]
|
|
)
|
|
|
|
resnets = []
|
|
for _ in range(num_layers):
|
|
resnets.append(
|
|
LTX2VideoResnetBlock3d(
|
|
in_channels=out_channels,
|
|
out_channels=out_channels,
|
|
dropout=dropout,
|
|
eps=resnet_eps,
|
|
non_linearity=resnet_act_fn,
|
|
inject_noise=inject_noise,
|
|
timestep_conditioning=timestep_conditioning,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
generator: torch.Generator | None = None,
|
|
causal: bool = True,
|
|
feat_cache: list | None = None,
|
|
feat_idx: list | None = None,
|
|
) -> torch.Tensor:
|
|
if self.conv_in is not None:
|
|
hidden_states = self.conv_in(
|
|
hidden_states, temb, generator, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx
|
|
)
|
|
|
|
if self.time_embedder is not None:
|
|
temb = self.time_embedder(
|
|
timestep=temb.flatten(),
|
|
resolution=None,
|
|
aspect_ratio=None,
|
|
batch_size=hidden_states.size(0),
|
|
hidden_dtype=hidden_states.dtype,
|
|
)
|
|
temb = temb.view(hidden_states.size(0), -1, 1, 1, 1)
|
|
|
|
if self.upsamplers is not None:
|
|
for upsampler in self.upsamplers:
|
|
hidden_states = upsampler(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
for i, resnet in enumerate(self.resnets):
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing and feat_cache is None:
|
|
hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb, generator, causal)
|
|
else:
|
|
hidden_states = resnet(
|
|
hidden_states, temb, generator, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx
|
|
)
|
|
|
|
return hidden_states
|
|
|
|
|
|
# =============================================================================
|
|
# Encoder
|
|
# =============================================================================
|
|
|
|
|
|
class LTX2VideoEncoder3d(nn.Module):
|
|
"""Encoder layer for encoding video samples to latent representation.
|
|
|
|
Args:
|
|
in_channels: Number of input channels.
|
|
out_channels: Number of latent channels.
|
|
block_out_channels: Number of output channels for each block.
|
|
spatio_temporal_scaling: Whether each block should contain downscaling.
|
|
layers_per_block: Number of layers per block.
|
|
downsample_type: Downsampling pattern per block.
|
|
patch_size: Size of spatial patches.
|
|
patch_size_t: Size of temporal patches.
|
|
resnet_norm_eps: Epsilon for ResNet normalization.
|
|
is_causal: Whether to use causal behavior.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int = 3,
|
|
out_channels: int = 128,
|
|
block_out_channels: tuple[int, ...] = (256, 512, 1024, 2048),
|
|
down_block_types: tuple[str, ...] = (
|
|
"LTX2VideoDownBlock3D",
|
|
"LTX2VideoDownBlock3D",
|
|
"LTX2VideoDownBlock3D",
|
|
"LTX2VideoDownBlock3D",
|
|
),
|
|
spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, True),
|
|
layers_per_block: tuple[int, ...] = (4, 6, 6, 2, 2),
|
|
downsample_type: tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
|
|
patch_size: int = 4,
|
|
patch_size_t: int = 1,
|
|
resnet_norm_eps: float = 1e-6,
|
|
is_causal: bool = True,
|
|
spatial_padding_mode: str = "zeros",
|
|
):
|
|
super().__init__()
|
|
|
|
self.patch_size = patch_size
|
|
self.patch_size_t = patch_size_t
|
|
self.in_channels = in_channels * patch_size**2 * patch_size_t
|
|
self.is_causal = is_causal
|
|
|
|
output_channel = out_channels
|
|
|
|
self.conv_in = LTX2VideoCausalConv3d(
|
|
in_channels=self.in_channels,
|
|
out_channels=output_channel,
|
|
kernel_size=3,
|
|
stride=1,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
# down blocks
|
|
num_block_out_channels = len(block_out_channels)
|
|
self.down_blocks = nn.ModuleList([])
|
|
for i in range(num_block_out_channels):
|
|
input_channel = output_channel
|
|
output_channel = block_out_channels[i]
|
|
|
|
if down_block_types[i] == "LTX2VideoDownBlock3D":
|
|
down_block = LTX2VideoDownBlock3D(
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
num_layers=layers_per_block[i],
|
|
resnet_eps=resnet_norm_eps,
|
|
spatio_temporal_scale=spatio_temporal_scaling[i],
|
|
downsample_type=downsample_type[i],
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown down block type: {down_block_types[i]}")
|
|
|
|
self.down_blocks.append(down_block)
|
|
|
|
# mid block
|
|
self.mid_block = LTX2VideoMidBlock3d(
|
|
in_channels=output_channel,
|
|
num_layers=layers_per_block[-1],
|
|
resnet_eps=resnet_norm_eps,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
# out
|
|
self.norm_out = PerChannelRMSNorm()
|
|
self.conv_act = nn.SiLU()
|
|
self.conv_out = LTX2VideoCausalConv3d(
|
|
in_channels=output_channel,
|
|
out_channels=out_channels + 1,
|
|
kernel_size=3,
|
|
stride=1,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
causal: bool | None = None,
|
|
feat_cache: list | None = None,
|
|
feat_idx: list | None = None,
|
|
) -> torch.Tensor:
|
|
p = self.patch_size
|
|
p_t = self.patch_size_t
|
|
|
|
batch_size, num_channels, num_frames, height, width = hidden_states.shape
|
|
post_patch_num_frames = num_frames // p_t
|
|
post_patch_height = height // p
|
|
post_patch_width = width // p
|
|
causal = causal or self.is_causal
|
|
if feat_cache is not None and feat_idx is None:
|
|
feat_idx = [0]
|
|
|
|
hidden_states = hidden_states.reshape(
|
|
batch_size, num_channels, post_patch_num_frames, p_t, post_patch_height, p, post_patch_width, p
|
|
)
|
|
hidden_states = hidden_states.permute(0, 1, 3, 7, 5, 2, 4, 6).flatten(1, 4)
|
|
hidden_states = self.conv_in(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing and feat_cache is None:
|
|
for down_block in self.down_blocks:
|
|
hidden_states = self._gradient_checkpointing_func(down_block, hidden_states, None, None, causal)
|
|
|
|
hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states, None, None, causal)
|
|
else:
|
|
for down_block in self.down_blocks:
|
|
hidden_states = down_block(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
hidden_states = self.mid_block(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
hidden_states = self.norm_out(hidden_states)
|
|
hidden_states = self.conv_act(hidden_states)
|
|
hidden_states = self.conv_out(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
last_channel = hidden_states[:, -1:]
|
|
last_channel = last_channel.repeat(1, hidden_states.size(1) - 2, 1, 1, 1)
|
|
hidden_states = torch.cat([hidden_states, last_channel], dim=1)
|
|
|
|
return hidden_states
|
|
|
|
|
|
# =============================================================================
|
|
# Decoder
|
|
# =============================================================================
|
|
|
|
|
|
class LTX2VideoDecoder3d(nn.Module):
|
|
"""Decoder layer for decoding latent representation to video.
|
|
|
|
Args:
|
|
in_channels: Number of latent channels.
|
|
out_channels: Number of output channels.
|
|
block_out_channels: Number of output channels for each block.
|
|
spatio_temporal_scaling: Whether each block should contain upscaling.
|
|
layers_per_block: Number of layers per block.
|
|
patch_size: Size of spatial patches.
|
|
patch_size_t: Size of temporal patches.
|
|
resnet_norm_eps: Epsilon for ResNet normalization.
|
|
is_causal: Whether to use causal behavior.
|
|
"""
|
|
|
|
_UPSAMPLE_STRIDE_MAP = {
|
|
"spatial": (1, 2, 2),
|
|
"temporal": (2, 1, 1),
|
|
"spatiotemporal": (2, 2, 2),
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int = 128,
|
|
out_channels: int = 3,
|
|
block_out_channels: tuple[int, ...] = (256, 512, 1024),
|
|
spatio_temporal_scaling: tuple[bool, ...] = (True, True, True),
|
|
layers_per_block: tuple[int, ...] = (5, 5, 5, 5),
|
|
patch_size: int = 4,
|
|
patch_size_t: int = 1,
|
|
resnet_norm_eps: float = 1e-6,
|
|
is_causal: bool = False,
|
|
inject_noise: tuple[bool, ...] = (False, False, False),
|
|
timestep_conditioning: bool = False,
|
|
upsample_residual: tuple[bool, ...] = (True, True, True),
|
|
upsample_factor: tuple[bool, ...] = (2, 2, 2),
|
|
upsample_type: tuple[str, ...] | None = None,
|
|
spatial_padding_mode: str = "reflect",
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self.patch_size = patch_size
|
|
self.patch_size_t = patch_size_t
|
|
self.out_channels = out_channels * patch_size**2
|
|
self.is_causal = is_causal
|
|
|
|
block_out_channels = tuple(reversed(block_out_channels))
|
|
spatio_temporal_scaling = tuple(reversed(spatio_temporal_scaling))
|
|
layers_per_block = tuple(reversed(layers_per_block))
|
|
inject_noise = tuple(reversed(inject_noise))
|
|
upsample_residual = tuple(reversed(upsample_residual))
|
|
upsample_factor = tuple(reversed(upsample_factor))
|
|
if upsample_type is not None:
|
|
upsample_type = tuple(reversed(upsample_type))
|
|
output_channel = block_out_channels[0]
|
|
|
|
self.conv_in = LTX2VideoCausalConv3d(
|
|
in_channels=in_channels,
|
|
out_channels=output_channel,
|
|
kernel_size=3,
|
|
stride=1,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
self.mid_block = LTX2VideoMidBlock3d(
|
|
in_channels=output_channel,
|
|
num_layers=layers_per_block[0],
|
|
resnet_eps=resnet_norm_eps,
|
|
inject_noise=inject_noise[0],
|
|
timestep_conditioning=timestep_conditioning,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
# up blocks
|
|
num_block_out_channels = len(block_out_channels)
|
|
self.up_blocks = nn.ModuleList([])
|
|
for i in range(num_block_out_channels):
|
|
input_channel = output_channel // upsample_factor[i]
|
|
output_channel = block_out_channels[i] // upsample_factor[i]
|
|
|
|
stride = (2, 2, 2)
|
|
if upsample_type is not None:
|
|
stride = self._UPSAMPLE_STRIDE_MAP[upsample_type[i]]
|
|
|
|
up_block = LTX2VideoUpBlock3d(
|
|
in_channels=input_channel,
|
|
out_channels=output_channel,
|
|
num_layers=layers_per_block[i + 1],
|
|
resnet_eps=resnet_norm_eps,
|
|
spatio_temporal_scale=spatio_temporal_scaling[i],
|
|
inject_noise=inject_noise[i + 1],
|
|
timestep_conditioning=timestep_conditioning,
|
|
upsample_residual=upsample_residual[i],
|
|
upscale_factor=upsample_factor[i],
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
upsample_stride=stride,
|
|
)
|
|
|
|
self.up_blocks.append(up_block)
|
|
|
|
# out
|
|
self.norm_out = PerChannelRMSNorm()
|
|
self.conv_act = nn.SiLU()
|
|
self.conv_out = LTX2VideoCausalConv3d(
|
|
in_channels=output_channel,
|
|
out_channels=self.out_channels,
|
|
kernel_size=3,
|
|
stride=1,
|
|
spatial_padding_mode=spatial_padding_mode,
|
|
)
|
|
|
|
# timestep embedding
|
|
self.time_embedder = None
|
|
self.scale_shift_table = None
|
|
self.timestep_scale_multiplier = None
|
|
if timestep_conditioning:
|
|
self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0, dtype=torch.float32))
|
|
self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(output_channel * 2, 0)
|
|
self.scale_shift_table = nn.Parameter(torch.randn(2, output_channel) / output_channel**0.5)
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
causal: bool | None = None,
|
|
feat_cache: list | None = None,
|
|
feat_idx: list | None = None,
|
|
) -> torch.Tensor:
|
|
causal = causal or self.is_causal
|
|
_shape_of(hidden_states)
|
|
|
|
if feat_cache is not None and feat_idx is None:
|
|
feat_idx = [0]
|
|
|
|
hidden_states = self.conv_in(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
if self.timestep_scale_multiplier is not None:
|
|
temb = temb * self.timestep_scale_multiplier
|
|
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing and feat_cache is None:
|
|
hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states, temb, None, causal)
|
|
|
|
for up_block in self.up_blocks:
|
|
hidden_states = self._gradient_checkpointing_func(up_block, hidden_states, temb, None, causal)
|
|
else:
|
|
hidden_states = self.mid_block(hidden_states, temb, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
for up_block in self.up_blocks:
|
|
hidden_states = up_block(hidden_states, temb, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
hidden_states = self.norm_out(hidden_states)
|
|
|
|
if self.time_embedder is not None:
|
|
temb = self.time_embedder(
|
|
timestep=temb.flatten(),
|
|
resolution=None,
|
|
aspect_ratio=None,
|
|
batch_size=hidden_states.size(0),
|
|
hidden_dtype=hidden_states.dtype,
|
|
)
|
|
temb = temb.view(hidden_states.size(0), -1, 1, 1, 1).unflatten(1, (2, -1))
|
|
temb = temb + self.scale_shift_table[None, ..., None, None, None]
|
|
shift, scale = temb.unbind(dim=1)
|
|
hidden_states = hidden_states * (1 + scale) + shift
|
|
|
|
hidden_states = self.conv_act(hidden_states)
|
|
hidden_states = self.conv_out(hidden_states, causal=causal, feat_cache=feat_cache, feat_idx=feat_idx)
|
|
|
|
p = self.patch_size
|
|
p_t = self.patch_size_t
|
|
|
|
batch_size, num_channels, num_frames, height, width = hidden_states.shape
|
|
hidden_states = hidden_states.reshape(batch_size, -1, p_t, p, p, num_frames, height, width)
|
|
hidden_states = hidden_states.permute(0, 1, 5, 2, 6, 4, 7, 3).flatten(6, 7).flatten(4, 5).flatten(2, 3)
|
|
|
|
return hidden_states
|
|
|
|
|
|
# =============================================================================
|
|
# Main VAE Model
|
|
# =============================================================================
|
|
|
|
|
|
class AutoencoderKLCausalLTX2Video(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin):
|
|
"""A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
|
|
|
|
Used in LTX-2. Supports causal encoding/decoding with streaming cache for memory-efficient
|
|
processing of long videos.
|
|
|
|
Args:
|
|
in_channels: Number of input channels.
|
|
out_channels: Number of output channels.
|
|
latent_channels: Number of latent channels.
|
|
block_out_channels: Number of output channels for each encoder block.
|
|
decoder_block_out_channels: Number of output channels for each decoder block.
|
|
layers_per_block: Number of layers per encoder block.
|
|
decoder_layers_per_block: Number of layers per decoder block.
|
|
spatio_temporal_scaling: Whether each encoder block should downscale.
|
|
decoder_spatio_temporal_scaling: Whether each decoder block should upscale.
|
|
downsample_type: Downsampling pattern for encoder.
|
|
upsample_residual: Whether to use residual in decoder upsampling.
|
|
upsample_factor: Upsampling factor for each decoder block.
|
|
timestep_conditioning: Whether to condition on timesteps.
|
|
patch_size: Size of spatial patches.
|
|
patch_size_t: Size of temporal patches.
|
|
resnet_norm_eps: Epsilon for ResNet normalization.
|
|
scaling_factor: Factor to scale latents.
|
|
encoder_causal: Whether encoder should be causal.
|
|
decoder_causal: Whether decoder should be causal.
|
|
gradient_checkpointing: Whether to use gradient checkpointing.
|
|
"""
|
|
|
|
_supports_gradient_checkpointing = True
|
|
|
|
@register_to_config
|
|
def __init__(
|
|
self,
|
|
in_channels: int = 3,
|
|
out_channels: int = 3,
|
|
latent_channels: int = 128,
|
|
block_out_channels: tuple[int, ...] = (256, 512, 1024, 2048),
|
|
down_block_types: tuple[str, ...] = (
|
|
"LTX2VideoDownBlock3D",
|
|
"LTX2VideoDownBlock3D",
|
|
"LTX2VideoDownBlock3D",
|
|
"LTX2VideoDownBlock3D",
|
|
),
|
|
decoder_block_out_channels: tuple[int, ...] = (256, 512, 1024),
|
|
layers_per_block: tuple[int, ...] = (4, 6, 6, 2, 2),
|
|
decoder_layers_per_block: tuple[int, ...] = (5, 5, 5, 5),
|
|
spatio_temporal_scaling: tuple[bool, ...] = (True, True, True, True),
|
|
decoder_spatio_temporal_scaling: tuple[bool, ...] = (True, True, True),
|
|
decoder_inject_noise: tuple[bool, ...] = (False, False, False, False),
|
|
downsample_type: tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
|
|
upsample_residual: tuple[bool, ...] = (True, True, True),
|
|
upsample_factor: tuple[int, ...] = (2, 2, 2),
|
|
decoder_upsample_type: tuple[str, ...] | None = None,
|
|
timestep_conditioning: bool = False,
|
|
patch_size: int = 4,
|
|
patch_size_t: int = 1,
|
|
resnet_norm_eps: float = 1e-6,
|
|
scaling_factor: float = 1.0,
|
|
encoder_causal: bool = True,
|
|
decoder_causal: bool = True,
|
|
gradient_checkpointing: bool = False,
|
|
encoder_spatial_padding_mode: str = "zeros",
|
|
decoder_spatial_padding_mode: str = "reflect",
|
|
spatial_compression_ratio: int = None,
|
|
temporal_compression_ratio: int = None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.encoder = LTX2VideoEncoder3d(
|
|
in_channels=in_channels,
|
|
out_channels=latent_channels,
|
|
block_out_channels=block_out_channels,
|
|
down_block_types=down_block_types,
|
|
spatio_temporal_scaling=spatio_temporal_scaling,
|
|
layers_per_block=layers_per_block,
|
|
downsample_type=downsample_type,
|
|
patch_size=patch_size,
|
|
patch_size_t=patch_size_t,
|
|
resnet_norm_eps=resnet_norm_eps,
|
|
is_causal=encoder_causal,
|
|
spatial_padding_mode=encoder_spatial_padding_mode,
|
|
)
|
|
self.decoder = LTX2VideoDecoder3d(
|
|
in_channels=latent_channels,
|
|
out_channels=out_channels,
|
|
block_out_channels=decoder_block_out_channels,
|
|
spatio_temporal_scaling=decoder_spatio_temporal_scaling,
|
|
layers_per_block=decoder_layers_per_block,
|
|
patch_size=patch_size,
|
|
patch_size_t=patch_size_t,
|
|
resnet_norm_eps=resnet_norm_eps,
|
|
is_causal=decoder_causal,
|
|
timestep_conditioning=timestep_conditioning,
|
|
inject_noise=decoder_inject_noise,
|
|
upsample_residual=upsample_residual,
|
|
upsample_factor=upsample_factor,
|
|
upsample_type=decoder_upsample_type,
|
|
spatial_padding_mode=decoder_spatial_padding_mode,
|
|
)
|
|
|
|
latents_mean = torch.zeros((latent_channels,), requires_grad=False)
|
|
latents_std = torch.ones((latent_channels,), requires_grad=False)
|
|
self.register_buffer("latents_mean", latents_mean, persistent=True)
|
|
self.register_buffer("latents_std", latents_std, persistent=True)
|
|
|
|
self.spatial_compression_ratio = (
|
|
patch_size * 2 ** sum(spatio_temporal_scaling)
|
|
if spatial_compression_ratio is None
|
|
else spatial_compression_ratio
|
|
)
|
|
self.temporal_compression_ratio = (
|
|
patch_size_t * 2 ** sum(spatio_temporal_scaling)
|
|
if temporal_compression_ratio is None
|
|
else temporal_compression_ratio
|
|
)
|
|
|
|
# Memory optimization flags
|
|
self.use_slicing = False
|
|
self.use_tiling = False
|
|
self.use_framewise_encoding = False
|
|
self.use_framewise_decoding = False
|
|
|
|
# Tiling configuration
|
|
self.num_sample_frames_batch_size = 16
|
|
self.num_latent_frames_batch_size = 2
|
|
self.tile_sample_min_height = 512
|
|
self.tile_sample_min_width = 512
|
|
self.tile_sample_min_num_frames = 16
|
|
self.tile_sample_stride_height = 448
|
|
self.tile_sample_stride_width = 448
|
|
self.tile_sample_stride_num_frames = 8
|
|
|
|
# Cache managers for streaming inference
|
|
self._encoder_cache = EncoderCacheManager()
|
|
self._decoder_cache = DecoderCacheManager()
|
|
|
|
# Expose config-driven gradient checkpointing toggle for training
|
|
self.gradient_checkpointing = gradient_checkpointing
|
|
if gradient_checkpointing:
|
|
self.enable_gradient_checkpointing()
|
|
|
|
def enable_tiling(
|
|
self,
|
|
tile_sample_min_height: int | None = None,
|
|
tile_sample_min_width: int | None = None,
|
|
tile_sample_min_num_frames: int | None = None,
|
|
tile_sample_stride_height: float | None = None,
|
|
tile_sample_stride_width: float | None = None,
|
|
tile_sample_stride_num_frames: float | None = None,
|
|
) -> None:
|
|
"""Enable tiled VAE decoding for memory-efficient processing of large videos."""
|
|
self.use_tiling = True
|
|
self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
|
|
self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
|
|
self.tile_sample_min_num_frames = tile_sample_min_num_frames or self.tile_sample_min_num_frames
|
|
self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
|
|
self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
|
|
self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames
|
|
|
|
def _encode(self, x: torch.Tensor, causal: bool | None = None) -> torch.Tensor:
|
|
"""Internal encode method."""
|
|
batch_size, num_channels, num_frames, height, width = x.shape
|
|
|
|
if self.use_framewise_decoding and num_frames > self.tile_sample_min_num_frames:
|
|
return self._temporal_tiled_encode(x, causal=causal)
|
|
|
|
if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height):
|
|
return self.tiled_encode(x, causal=causal)
|
|
|
|
enc = self.encoder(x, causal=causal)
|
|
return enc
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Cache Management (Legacy API for backward compatibility)
|
|
# -------------------------------------------------------------------------
|
|
def clear_encoder_cache(self) -> None:
|
|
"""Clear encoder cache (legacy API, delegates to cache manager)."""
|
|
self._encoder_cache.clear()
|
|
|
|
def clear_decoder_cache(self) -> None:
|
|
"""Clear decoder cache (legacy API, delegates to cache manager)."""
|
|
self._decoder_cache.clear()
|
|
|
|
@property
|
|
def _encoder_feat_map(self) -> list:
|
|
"""Legacy property for encoder feature map."""
|
|
return self._encoder_cache.feat_map
|
|
|
|
@_encoder_feat_map.setter
|
|
def _encoder_feat_map(self, value: list):
|
|
"""Setter for legacy encoder feature map property."""
|
|
self._encoder_cache._feat_map = value
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Streaming Encode/Decode Methods
|
|
# -------------------------------------------------------------------------
|
|
def decode_with_cache(
|
|
self,
|
|
z: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
causal: bool | None = None,
|
|
return_dict: bool = True,
|
|
reset_cache: bool = False,
|
|
prepend_prev_latent_frames: int = 1,
|
|
) -> DecoderOutput | tuple[torch.Tensor]:
|
|
"""Decode one latent chunk with persistent decoder feature cache.
|
|
|
|
For decoder_causal=False model: decodes each chunk independently with:
|
|
- Left context from previous chunk
|
|
- Zero-appended right context placeholder
|
|
|
|
Args:
|
|
z: Input latent tensor [B, C, T, H, W]
|
|
temb: Optional timestep embedding
|
|
causal: Whether to use causal decoding
|
|
return_dict: Whether to return dict or tuple
|
|
reset_cache: Whether to reset cache before decoding
|
|
prepend_prev_latent_frames: Number of previous frames to prepend as context
|
|
|
|
Returns:
|
|
DecoderOutput containing decoded video
|
|
"""
|
|
if reset_cache:
|
|
self._decoder_cache.clear()
|
|
|
|
if self.use_slicing and z.shape[0] > 1:
|
|
raise ValueError("decode_with_cache does not support batch slicing; pass batch size 1 or disable slicing.")
|
|
if z.shape[2] <= 0:
|
|
raise ValueError("Input latent chunk must contain at least 1 frame.")
|
|
if prepend_prev_latent_frames < 0:
|
|
raise ValueError(f"`prepend_prev_latent_frames` must be >= 0, got {prepend_prev_latent_frames}.")
|
|
|
|
causal = self.decoder.is_causal if causal is None else causal
|
|
self._decoder_cache.validate_mode(causal)
|
|
|
|
ratio = self.temporal_compression_ratio
|
|
z_chunk = z
|
|
|
|
if not causal:
|
|
z_chunk = self._decoder_cache.prepend_context(z, prepend_prev_latent_frames, ratio)
|
|
|
|
decoded_full = self.decoder(
|
|
z_chunk,
|
|
temb,
|
|
causal=causal,
|
|
feat_cache=self._decoder_cache.feat_map,
|
|
feat_idx=[0],
|
|
)
|
|
|
|
if not causal:
|
|
decoded = self._decoder_cache.trim_output(decoded_full, prepend_prev_latent_frames, z.shape[2], ratio)
|
|
self._decoder_cache.update_tail(z, prepend_prev_latent_frames)
|
|
else:
|
|
decoded = decoded_full
|
|
|
|
if not return_dict:
|
|
return (decoded,)
|
|
return DecoderOutput(sample=decoded)
|
|
|
|
@apply_forward_hook
|
|
def decode_per_frame_with_cache(
|
|
self,
|
|
z: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
causal: bool | None = None,
|
|
reset_cache: bool = False,
|
|
) -> Iterator[torch.Tensor]:
|
|
"""Stream-decode a latent video one latent frame at a time, yielding pixel chunks.
|
|
|
|
Each iteration decodes a single latent frame and relies on the per-layer
|
|
feature cache for temporal consistency. Yields the decoded pixel tensor
|
|
for each latent frame so the caller can process/save/display each chunk
|
|
immediately (low peak memory, low latency, no final ``torch.cat``).
|
|
|
|
Streaming use: pass ``reset_cache=True`` on the first call of a new
|
|
stream/session, then ``reset_cache=False`` on subsequent calls so the
|
|
per-layer feature cache carries temporal state across calls.
|
|
|
|
Args:
|
|
z: Input latent tensor [B, C, T, H, W].
|
|
temb: Optional timestep embedding.
|
|
causal: Whether to use causal decoding.
|
|
reset_cache: If True, clear the decoder feature cache before decoding.
|
|
Set True only on the first chunk of a new stream.
|
|
|
|
Yields:
|
|
Decoded pixel tensor [B, C, T_out, H, W] for each single-latent-frame chunk.
|
|
"""
|
|
if self.use_slicing and z.shape[0] > 1:
|
|
raise ValueError(
|
|
"decode_per_frame_with_cache does not support batch slicing; pass batch size 1 or disable slicing."
|
|
)
|
|
if z.shape[2] <= 0:
|
|
raise ValueError("Input latent video must contain at least 1 frame.")
|
|
|
|
causal = self.decoder.is_causal if causal is None else causal
|
|
num_latent_frames = z.shape[2]
|
|
|
|
if reset_cache:
|
|
self._decoder_cache.clear()
|
|
self._decoder_cache.validate_mode(causal)
|
|
|
|
for t in range(num_latent_frames):
|
|
z_chunk = z[:, :, t : t + 1, :, :]
|
|
decoded_chunk = self.decoder(
|
|
z_chunk,
|
|
temb,
|
|
causal=causal,
|
|
feat_cache=self._decoder_cache.feat_map,
|
|
feat_idx=[0],
|
|
)
|
|
yield decoded_chunk
|
|
|
|
@apply_forward_hook
|
|
def streaming_causal_encode(
|
|
self,
|
|
x: torch.Tensor,
|
|
causal: bool | None = True,
|
|
return_dict: bool = True,
|
|
reset_cache: bool = False,
|
|
) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
|
|
"""Encode one streaming video chunk with persistent causal encoder cache.
|
|
|
|
Unlike ``encode``, this method does not split the input chunk and does
|
|
not clear encoder cache unless ``reset_cache`` is True.
|
|
"""
|
|
if reset_cache:
|
|
self._encoder_cache.clear()
|
|
|
|
if self.use_slicing and x.shape[0] > 1:
|
|
raise ValueError(
|
|
"streaming_causal_encode does not support batch slicing; pass batch size 1 or disable slicing."
|
|
)
|
|
if x.shape[2] <= 0:
|
|
raise ValueError("Input video chunk must contain at least 1 frame.")
|
|
|
|
h = self.encoder(x, causal=causal, feat_cache=self._encoder_cache.feat_map, feat_idx=[0])
|
|
posterior = DiagonalGaussianDistribution(h)
|
|
|
|
if not return_dict:
|
|
return (posterior,)
|
|
return AutoencoderKLOutput(latent_dist=posterior)
|
|
|
|
@apply_forward_hook
|
|
def encode(
|
|
self,
|
|
x: torch.Tensor,
|
|
causal: bool | None = True,
|
|
return_dict: bool = True,
|
|
) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]:
|
|
"""Encode a full video by iterating over temporal chunks with persistent encoder cache.
|
|
|
|
Args:
|
|
x: Input video tensor [B, C, T, H, W]
|
|
causal: Whether to use causal encoding.
|
|
return_dict: Whether to return dict or tuple.
|
|
|
|
Returns:
|
|
AutoencoderKLOutput containing latent distribution.
|
|
"""
|
|
chunk_num_frames = self.temporal_compression_ratio
|
|
num_frames = x.shape[2]
|
|
first_chunk_num_frames = num_frames % self.temporal_compression_ratio
|
|
|
|
if self.use_slicing and x.shape[0] > 1:
|
|
raise ValueError(
|
|
"encode_full_video_with_cache does not support batch slicing; pass batch size 1 or disable slicing."
|
|
)
|
|
|
|
self._encoder_cache.clear()
|
|
|
|
encoded_chunks = []
|
|
start = 0
|
|
step = first_chunk_num_frames if first_chunk_num_frames > 0 else chunk_num_frames
|
|
while start < num_frames:
|
|
end = min(start + step, num_frames)
|
|
encoded_chunk = self.encoder(
|
|
x[:, :, start:end, :, :],
|
|
causal=causal,
|
|
feat_cache=self._encoder_cache.feat_map,
|
|
feat_idx=[0],
|
|
)
|
|
encoded_chunks.append(encoded_chunk)
|
|
start = end
|
|
step = chunk_num_frames
|
|
|
|
if len(encoded_chunks) == 1:
|
|
h = encoded_chunks[0]
|
|
else:
|
|
h = torch.cat(encoded_chunks, dim=2)
|
|
|
|
self._encoder_cache.check_pending_consumed()
|
|
posterior = DiagonalGaussianDistribution(h)
|
|
|
|
if not return_dict:
|
|
return (posterior,)
|
|
return AutoencoderKLOutput(latent_dist=posterior)
|
|
|
|
def _decode(
|
|
self,
|
|
z: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
causal: bool | None = None,
|
|
return_dict: bool = True,
|
|
) -> DecoderOutput | torch.Tensor:
|
|
"""Internal decode method."""
|
|
batch_size, num_channels, num_frames, height, width = z.shape
|
|
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
|
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
|
tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
|
|
|
|
if self.use_framewise_decoding and num_frames > tile_latent_min_num_frames:
|
|
return self._temporal_tiled_decode(z, temb, causal=causal, return_dict=return_dict)
|
|
|
|
if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
|
|
return self.tiled_decode(z, temb, causal=causal, return_dict=return_dict)
|
|
|
|
dec = self.decoder(z, temb, causal=causal)
|
|
|
|
if not return_dict:
|
|
return (dec,)
|
|
|
|
return DecoderOutput(sample=dec)
|
|
|
|
@apply_forward_hook
|
|
def decode(
|
|
self,
|
|
z: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
causal: bool | None = None,
|
|
return_dict: bool = True,
|
|
) -> DecoderOutput | torch.Tensor:
|
|
"""Decode a batch of latents into images.
|
|
|
|
Args:
|
|
z: Input batch of latent vectors [B, C, T, H, W].
|
|
temb: Optional timestep embedding.
|
|
causal: Whether to use causal decoding.
|
|
return_dict: Whether to return a dict or tuple.
|
|
|
|
Returns:
|
|
DecoderOutput containing decoded video.
|
|
"""
|
|
if self.use_slicing and z.shape[0] > 1:
|
|
if temb is not None:
|
|
decoded_slices = [
|
|
self._decode(z_slice, t_slice, causal=causal).sample
|
|
for z_slice, t_slice in (z.split(1), temb.split(1))
|
|
]
|
|
else:
|
|
decoded_slices = [self._decode(z_slice, causal=causal).sample for z_slice in z.split(1)]
|
|
decoded = torch.cat(decoded_slices)
|
|
else:
|
|
decoded = self._decode(z, temb, causal=causal).sample
|
|
|
|
if not return_dict:
|
|
return (decoded,)
|
|
|
|
return DecoderOutput(sample=decoded)
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Tiling Methods
|
|
# -------------------------------------------------------------------------
|
|
|
|
def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
|
|
"""Blend vertically between two tiles."""
|
|
blend_extent = min(a.shape[3], b.shape[3], blend_extent)
|
|
for y in range(blend_extent):
|
|
b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * (
|
|
y / blend_extent
|
|
)
|
|
return b
|
|
|
|
def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
|
|
"""Blend horizontally between two tiles."""
|
|
blend_extent = min(a.shape[4], b.shape[4], blend_extent)
|
|
for x in range(blend_extent):
|
|
b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * (
|
|
x / blend_extent
|
|
)
|
|
return b
|
|
|
|
def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor:
|
|
"""Blend temporally between two tiles."""
|
|
blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
|
|
for x in range(blend_extent):
|
|
b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * (
|
|
x / blend_extent
|
|
)
|
|
return b
|
|
|
|
@staticmethod
|
|
def _get_tile_positions(total_size: int, tile_size: int, overlap: int) -> list:
|
|
"""Compute tile start positions ensuring full coverage with given overlap.
|
|
|
|
Returns a list of start positions such that every pixel in [0, total_size)
|
|
is covered by at least one tile of size ``tile_size``.
|
|
"""
|
|
if total_size <= tile_size:
|
|
return [0]
|
|
stride = tile_size - overlap
|
|
positions = list(range(0, total_size - tile_size, stride))
|
|
last = total_size - tile_size
|
|
if not positions or positions[-1] < last:
|
|
positions.append(last)
|
|
return positions
|
|
|
|
@staticmethod
|
|
def _make_blend_mask(
|
|
h: int,
|
|
w: int,
|
|
overlap_h: int,
|
|
overlap_w: int,
|
|
is_top: bool,
|
|
is_bottom: bool,
|
|
is_left: bool,
|
|
is_right: bool,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
) -> torch.Tensor:
|
|
"""Create a 2D spatial blend mask with linear ramps in overlap regions.
|
|
|
|
Interior pixels get weight 1. Pixels in the overlap band of a non-edge
|
|
side get linearly increasing/decreasing weights so that the weighted
|
|
average of overlapping tiles transitions smoothly.
|
|
|
|
Returns:
|
|
Tensor of shape ``[h, w]`` with values in (0, 1].
|
|
"""
|
|
mask_h = torch.ones(h, device=device, dtype=dtype)
|
|
mask_w = torch.ones(w, device=device, dtype=dtype)
|
|
|
|
if not is_top and overlap_h > 0:
|
|
mask_h[:overlap_h] = torch.linspace(0, 1, overlap_h + 2, device=device, dtype=dtype)[1:-1]
|
|
if not is_bottom and overlap_h > 0:
|
|
mask_h[-overlap_h:] = torch.linspace(1, 0, overlap_h + 2, device=device, dtype=dtype)[1:-1]
|
|
if not is_left and overlap_w > 0:
|
|
mask_w[:overlap_w] = torch.linspace(0, 1, overlap_w + 2, device=device, dtype=dtype)[1:-1]
|
|
if not is_right and overlap_w > 0:
|
|
mask_w[-overlap_w:] = torch.linspace(1, 0, overlap_w + 2, device=device, dtype=dtype)[1:-1]
|
|
|
|
return mask_h[:, None] * mask_w[None, :]
|
|
|
|
def tiled_encode(
|
|
self,
|
|
x: torch.Tensor,
|
|
causal: bool | None = None,
|
|
tile_caches: dict | None = None,
|
|
) -> torch.Tensor:
|
|
r"""Encode a batch of images using a tiled encoder.
|
|
|
|
Args:
|
|
x: Input batch of videos [B, C, T, H, W].
|
|
causal: Whether to use causal encoding.
|
|
tile_caches: Optional dict mapping ``(row_idx, col_idx)`` to a
|
|
per-tile ``feat_cache`` list. When provided each tile's
|
|
encoder call receives its own persistent cache so temporal
|
|
context is maintained across successive calls.
|
|
|
|
Returns:
|
|
The latent representation of the encoded videos.
|
|
"""
|
|
batch_size, num_channels, num_frames, height, width = x.shape
|
|
latent_height = height // self.spatial_compression_ratio
|
|
latent_width = width // self.spatial_compression_ratio
|
|
|
|
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
|
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
|
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
|
|
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
|
|
|
|
blend_height = tile_latent_min_height - tile_latent_stride_height
|
|
blend_width = tile_latent_min_width - tile_latent_stride_width
|
|
|
|
# Split x into overlapping tiles and encode them separately.
|
|
# The tiles have an overlap to avoid seams between tiles.
|
|
rows = []
|
|
for i_idx, i in enumerate(range(0, height, self.tile_sample_stride_height)):
|
|
row = []
|
|
for j_idx, j in enumerate(range(0, width, self.tile_sample_stride_width)):
|
|
tile = x[:, :, :, i : i + self.tile_sample_min_height, j : j + self.tile_sample_min_width]
|
|
if tile_caches is not None:
|
|
enc_tile = self.encoder(
|
|
tile,
|
|
causal=causal,
|
|
feat_cache=tile_caches[(i_idx, j_idx)],
|
|
feat_idx=[0],
|
|
)
|
|
else:
|
|
enc_tile = self.encoder(tile, causal=causal)
|
|
row.append(enc_tile)
|
|
rows.append(row)
|
|
|
|
result_rows = []
|
|
for i, row in enumerate(rows):
|
|
result_row = []
|
|
for j, tile in enumerate(row):
|
|
# blend the above tile and the left tile
|
|
# to the current tile and add the current tile to the result row
|
|
if i > 0:
|
|
tile = self.blend_v(rows[i - 1][j], tile, blend_height)
|
|
if j > 0:
|
|
tile = self.blend_h(row[j - 1], tile, blend_width)
|
|
result_row.append(tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width])
|
|
result_rows.append(torch.cat(result_row, dim=4))
|
|
|
|
enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
|
|
return enc
|
|
|
|
def tiled_decode(
|
|
self,
|
|
z: torch.Tensor,
|
|
temb: torch.Tensor | None,
|
|
causal: bool | None = None,
|
|
return_dict: bool = True,
|
|
tile_caches: dict | None = None,
|
|
) -> DecoderOutput | torch.Tensor:
|
|
r"""Decode a batch of images using a tiled decoder.
|
|
|
|
Args:
|
|
z: Input batch of latent vectors [B, C, T, H, W].
|
|
temb: Optional timestep embedding.
|
|
causal: Whether to use causal decoding.
|
|
return_dict: Whether to return a DecoderOutput instead of a plain tuple.
|
|
tile_caches: Optional dict mapping ``(row_idx, col_idx)`` to a
|
|
per-tile ``feat_cache`` list. When provided each tile's
|
|
decoder call receives its own persistent cache so temporal
|
|
context is maintained across successive calls.
|
|
|
|
Returns:
|
|
DecoderOutput or tuple containing decoded video.
|
|
"""
|
|
|
|
batch_size, num_channels, num_frames, height, width = z.shape
|
|
sample_height = height * self.spatial_compression_ratio
|
|
sample_width = width * self.spatial_compression_ratio
|
|
|
|
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
|
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
|
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
|
|
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
|
|
|
|
blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
|
|
blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
|
|
|
|
# Split z into overlapping tiles and decode them separately.
|
|
# The tiles have an overlap to avoid seams between tiles.
|
|
rows = []
|
|
for i_idx, i in enumerate(range(0, height, tile_latent_stride_height)):
|
|
row = []
|
|
for j_idx, j in enumerate(range(0, width, tile_latent_stride_width)):
|
|
tile = z[:, :, :, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
|
|
if tile_caches is not None:
|
|
dec_tile = self.decoder(
|
|
tile,
|
|
temb,
|
|
causal=causal,
|
|
feat_cache=tile_caches[(i_idx, j_idx)],
|
|
feat_idx=[0],
|
|
)
|
|
else:
|
|
dec_tile = self.decoder(tile, temb, causal=causal)
|
|
row.append(dec_tile)
|
|
rows.append(row)
|
|
|
|
result_rows = []
|
|
for i, row in enumerate(rows):
|
|
result_row = []
|
|
for j, tile in enumerate(row):
|
|
# blend the above tile and the left tile
|
|
# to the current tile and add the current tile to the result row
|
|
if i > 0:
|
|
tile = self.blend_v(rows[i - 1][j], tile, blend_height)
|
|
if j > 0:
|
|
tile = self.blend_h(row[j - 1], tile, blend_width)
|
|
result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
|
|
result_rows.append(torch.cat(result_row, dim=4))
|
|
|
|
dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
|
|
|
|
if not return_dict:
|
|
return (dec,)
|
|
|
|
return DecoderOutput(sample=dec)
|
|
|
|
def decode_chunk_tile(
|
|
self, z: torch.Tensor, temb: torch.Tensor | None, causal: bool | None = None, return_dict: bool = True
|
|
) -> DecoderOutput | torch.Tensor:
|
|
"""Decode latent videos with temporal-only tiling.
|
|
|
|
This is used by SANA-Streaming inference with the public LTX-2 VAE
|
|
weights. It avoids spatial tiling and only chunks along latent time,
|
|
matching the release recipe used for long 720p V2V editing.
|
|
"""
|
|
batch_size, num_channels, num_frames, height, width = z.shape
|
|
del batch_size, num_channels, height, width
|
|
|
|
num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
|
|
tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
|
|
tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio
|
|
blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
|
|
|
|
row = []
|
|
for i in range(0, num_frames, tile_latent_stride_num_frames):
|
|
tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :]
|
|
decoded = self.decoder(tile, temb, causal=causal)
|
|
if i > 0:
|
|
decoded = decoded[:, :, 1:, :, :]
|
|
row.append(decoded)
|
|
|
|
result_row = []
|
|
for i, tile in enumerate(row):
|
|
if i > 0:
|
|
tile = self.blend_t(row[i - 1], tile, blend_num_frames)
|
|
tile = tile[:, :, : self.tile_sample_stride_num_frames, :, :]
|
|
result_row.append(tile)
|
|
else:
|
|
result_row.append(tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :])
|
|
|
|
dec = torch.cat(result_row, dim=2)[:, :, :num_sample_frames]
|
|
|
|
if not return_dict:
|
|
return (dec,)
|
|
return DecoderOutput(sample=dec)
|
|
|
|
def _temporal_tiled_encode(self, x: torch.Tensor, causal: bool | None = None) -> torch.Tensor:
|
|
"""Encode with temporal tiling for long videos."""
|
|
batch_size, num_channels, num_frames, height, width = x.shape
|
|
latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1
|
|
|
|
tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
|
|
tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio
|
|
blend_num_frames = tile_latent_min_num_frames - tile_latent_stride_num_frames
|
|
|
|
row = []
|
|
for i in range(0, num_frames, self.tile_sample_stride_num_frames):
|
|
tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :]
|
|
if self.use_tiling and (height > self.tile_sample_min_height or width > self.tile_sample_min_width):
|
|
tile = self.tiled_encode(tile, causal=causal)
|
|
else:
|
|
tile = self.encoder(tile, causal=causal)
|
|
if i > 0:
|
|
tile = tile[:, :, 1:, :, :]
|
|
row.append(tile)
|
|
|
|
result_row = []
|
|
for i, tile in enumerate(row):
|
|
if i > 0:
|
|
tile = self.blend_t(row[i - 1], tile, blend_num_frames)
|
|
result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :])
|
|
else:
|
|
result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :])
|
|
|
|
enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames]
|
|
return enc
|
|
|
|
def _temporal_tiled_decode(
|
|
self, z: torch.Tensor, temb: torch.Tensor | None, causal: bool | None = None, return_dict: bool = True
|
|
) -> DecoderOutput | torch.Tensor:
|
|
"""Decode with temporal tiling for long videos."""
|
|
batch_size, num_channels, num_frames, height, width = z.shape
|
|
num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
|
|
|
|
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
|
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
|
tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio
|
|
tile_latent_stride_num_frames = self.tile_sample_stride_num_frames // self.temporal_compression_ratio
|
|
blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
|
|
|
|
row = []
|
|
for i in range(0, num_frames, tile_latent_stride_num_frames):
|
|
tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :]
|
|
if self.use_tiling and (tile.shape[-1] > tile_latent_min_width or tile.shape[-2] > tile_latent_min_height):
|
|
decoded = self.tiled_decode(tile, temb, causal=causal, return_dict=True).sample
|
|
else:
|
|
decoded = self.decoder(tile, temb, causal=causal)
|
|
if i > 0:
|
|
decoded = decoded[:, :, 1:, :, :]
|
|
row.append(decoded)
|
|
|
|
result_row = []
|
|
for i, tile in enumerate(row):
|
|
if i > 0:
|
|
tile = self.blend_t(row[i - 1], tile, blend_num_frames)
|
|
tile = tile[:, :, : self.tile_sample_stride_num_frames, :, :]
|
|
result_row.append(tile)
|
|
else:
|
|
result_row.append(tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :])
|
|
|
|
dec = torch.cat(result_row, dim=2)[:, :, :num_sample_frames]
|
|
|
|
if not return_dict:
|
|
return (dec,)
|
|
return DecoderOutput(sample=dec)
|
|
|
|
def forward(
|
|
self,
|
|
sample: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
sample_posterior: bool = False,
|
|
encoder_causal: bool | None = None,
|
|
decoder_causal: bool | None = None,
|
|
return_dict: bool = True,
|
|
generator: torch.Generator | None = None,
|
|
) -> torch.Tensor | torch.Tensor:
|
|
"""Full forward pass: encode then decode.
|
|
|
|
Args:
|
|
sample: Input video tensor [B, C, T, H, W]
|
|
temb: Optional timestep embedding
|
|
sample_posterior: Whether to sample from posterior or use mode
|
|
encoder_causal: Whether to use causal encoding
|
|
decoder_causal: Whether to use causal decoding
|
|
return_dict: Whether to return dict
|
|
generator: Optional random generator
|
|
|
|
Returns:
|
|
Decoded video
|
|
"""
|
|
x = sample
|
|
posterior = self.encode(x, causal=encoder_causal).latent_dist
|
|
if sample_posterior:
|
|
z = posterior.sample(generator=generator)
|
|
else:
|
|
z = posterior.mode()
|
|
dec = self.decode(z, temb, causal=decoder_causal)
|
|
if not return_dict:
|
|
return (dec.sample,)
|
|
return dec
|