516 lines
22 KiB
Python
516 lines
22 KiB
Python
"""LTX 2.4 diffusion video VAE decoder (NADiffusionDecoder).
|
|
|
|
Port of the reference ``DiffusionVideoDecoder`` without the NATTEN dependency:
|
|
``natten.na3d`` is replaced by ``comfy_kitchen.na3d``, which reproduces
|
|
NATTEN's semantics (window of exactly ``kernel_size`` per query, shifted
|
|
inward at grid boundaries, dilation 1) and dispatches cuda/triton/eager per
|
|
device and dtype (the eager backend covers CPU and fp32).
|
|
|
|
Stages 1-4 deterministically upsample the latent into a context volume via
|
|
NA transformer blocks + linear pixel-shuffle upsamples. Stage 5 runs
|
|
``DiffusionNABlock``s that denoise patchified noised pixels ``x_t`` guided by
|
|
that context through AdaLN-Zero scale/shift. The 2.4 checkpoint is single-step
|
|
``x0``: one forward pass yields the pixels directly, no Euler loop.
|
|
|
|
State dict keys match the shipped checkpoints directly (fused ``attn.qkv``,
|
|
``t_embedder.mlp.{0,2}``, ``shared_adaln.proj``); no rename pass is needed.
|
|
"""
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from einops import rearrange
|
|
from torch import nn
|
|
|
|
from comfy.ldm.lightricks.model import get_timestep_embedding
|
|
from .causal_video_autoencoder import Encoder, processor
|
|
|
|
import comfy_kitchen
|
|
|
|
# Token chunk for the SwiGLU MLP (bounds the [chunk, hidden] workspace).
|
|
MLP_TOKEN_CHUNK = 65536
|
|
|
|
|
|
def rms_norm(x, weight, eps=1e-6):
|
|
if hasattr(F, "rms_norm"):
|
|
return F.rms_norm(x, (x.shape[-1],), weight=weight.to(x.dtype), eps=eps)
|
|
x_f = x.float()
|
|
x_f = x_f * torch.rsqrt(x_f.pow(2).mean(-1, keepdim=True) + eps)
|
|
return (x_f * weight.float()).to(x.dtype)
|
|
|
|
|
|
class RMSNorm(nn.Module):
|
|
def __init__(self, dim, eps=1e-6):
|
|
super().__init__()
|
|
self.eps = eps
|
|
self.weight = nn.Parameter(torch.ones(dim))
|
|
|
|
def forward(self, x):
|
|
return rms_norm(x, self.weight, self.eps)
|
|
|
|
|
|
def patchify(x, patch_size_hw, patch_size_t=1):
|
|
if patch_size_hw == 1 and patch_size_t == 1:
|
|
return x
|
|
return rearrange(x, "b c (f p) (h q) (w r) -> b (c p r q) f h w", p=patch_size_t, q=patch_size_hw, r=patch_size_hw)
|
|
|
|
|
|
def unpatchify(x, patch_size_hw, patch_size_t=1):
|
|
if patch_size_hw == 1 and patch_size_t == 1:
|
|
return x
|
|
return rearrange(x, "b (c p r q) f h w -> b c (f p) (h q) (w r)", p=patch_size_t, q=patch_size_hw, r=patch_size_hw)
|
|
|
|
|
|
# --- Absolute per-axis RoPE (matches ltx-core rope.py numerics) ---
|
|
|
|
def default_rope_dim_split(head_dim):
|
|
d_t = (head_dim // 4) // 2 * 2
|
|
d_hw = (head_dim - d_t) // 2
|
|
if d_hw % 2 != 0:
|
|
d_t -= 2
|
|
d_hw = (head_dim - d_t) // 2
|
|
return (d_t, d_hw, d_hw)
|
|
|
|
|
|
def rope_inv_freqs(dim, base=10000.0, device=None):
|
|
exponents = torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim
|
|
return (1.0 / torch.pow(torch.tensor(float(base), dtype=torch.float64, device=device), exponents)).to(torch.float32)
|
|
|
|
|
|
def _rope_tables(lengths, inv_freqs, device):
|
|
"""Precompute per-axis fp32 cos/sin tables for global 0-based positions."""
|
|
tables = []
|
|
for length, inv in zip(lengths, inv_freqs):
|
|
pos = torch.arange(length, dtype=torch.float32, device=device)
|
|
ang = pos[:, None] * inv[None, :]
|
|
tables.append((ang.cos(), ang.sin()))
|
|
return tables
|
|
|
|
|
|
def _rope_matrices_slice(tables, t0, t1, h, w):
|
|
"""Per-token rotation matrices ``(1, ts*h*w, 1, hd/2, 2, 2)`` fp32 for
|
|
``comfy_kitchen.rms_rope_`` (interleaved-pair convention), covering global
|
|
frames ``[t0, t1)`` of the axis-factorized tables."""
|
|
parts = []
|
|
for (c, s), sl in zip(tables, (slice(t0, t1), slice(None), slice(None))):
|
|
c, s = c[sl], s[sl]
|
|
parts.append(torch.stack([c, -s, s, c], dim=-1).reshape(c.shape[0], 1, 1, c.shape[1], 2, 2))
|
|
ts = t1 - t0
|
|
freqs = torch.cat([
|
|
parts[0].expand(ts, h, w, -1, 2, 2),
|
|
parts[1].transpose(0, 1).expand(ts, h, w, -1, 2, 2),
|
|
parts[2].movedim(0, 2).expand(ts, h, w, -1, 2, 2),
|
|
], dim=3)
|
|
return freqs.reshape(1, ts * h * w, 1, -1, 2, 2)
|
|
|
|
|
|
class NeighborhoodAttention3D(nn.Module):
|
|
"""QKV (fused, matching checkpoint keys) + q/k RMSNorm + abs RoPE + NA."""
|
|
|
|
def __init__(self, dim, kernel_size, head_dim=64, rope_base=10000.0):
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.num_heads = dim // head_dim
|
|
self.head_dim = head_dim
|
|
self.kernel_size = tuple(kernel_size)
|
|
self.scale = head_dim ** -0.5
|
|
self.rope_split = default_rope_dim_split(head_dim)
|
|
self.rope_base = rope_base
|
|
|
|
self.qkv = nn.Linear(dim, dim * 3, bias=True)
|
|
self.proj = nn.Linear(dim, dim, bias=True)
|
|
self.q_norm = RMSNorm(head_dim, eps=1e-6)
|
|
self.k_norm = RMSNorm(head_dim, eps=1e-6)
|
|
|
|
def forward(self, x, pre=None, add_to=None):
|
|
"""``pre`` (per-token norm/modulate) is applied slice-wise so the full
|
|
pre-attention tensor is never materialized; ``add_to`` streams the
|
|
output projection into it in place (residual add) and returns it.
|
|
Both bound peak memory without changing results."""
|
|
batch, t, h, w, _ = x.shape
|
|
inv_freqs = tuple(rope_inv_freqs(d, self.rope_base, device=x.device) for d in self.rope_split)
|
|
tables = _rope_tables((t, h, w), inv_freqs, x.device)
|
|
shape = (batch, t, h, w, self.num_heads, self.head_dim)
|
|
q = torch.empty(shape, dtype=x.dtype, device=x.device)
|
|
k = torch.empty(shape, dtype=x.dtype, device=x.device)
|
|
v = torch.empty(shape, dtype=x.dtype, device=x.device)
|
|
q_weight = (self.q_norm.weight.detach() * self.scale).to(x.dtype) # scale commutes with the rotation
|
|
k_weight = self.k_norm.weight.detach().to(x.dtype)
|
|
chunk = max(1, (2 ** 25) // max(h * w * self.dim, 1))
|
|
for t0 in range(0, t, chunk):
|
|
t1 = min(t0 + chunk, t)
|
|
sl = x[:, t0:t1] if pre is None else pre(x[:, t0:t1])
|
|
qc, kc, vc = self.qkv(sl).chunk(3, dim=-1)
|
|
cshape = (batch, t1 - t0, h, w, self.num_heads, self.head_dim)
|
|
q[:, t0:t1] = qc.reshape(cshape)
|
|
k[:, t0:t1] = kc.reshape(cshape)
|
|
v[:, t0:t1] = vc.reshape(cshape)
|
|
freqs = _rope_matrices_slice(tables, t0, t1, h, w)
|
|
nt = (t1 - t0) * h * w
|
|
for b in range(batch):
|
|
comfy_kitchen.rms_rope_(
|
|
q[b, t0:t1].view(1, nt, self.num_heads, self.head_dim),
|
|
k[b, t0:t1].view(1, nt, self.num_heads, self.head_dim),
|
|
freqs, q_weight, k_weight)
|
|
out = comfy_kitchen.na3d(q, k, v, list(self.kernel_size), None, 1.0)
|
|
del q, k, v
|
|
out = out.reshape(batch, t, h, w, self.dim)
|
|
res = add_to if add_to is not None else torch.empty_like(out)
|
|
for t0 in range(0, t, chunk):
|
|
t1 = min(t0 + chunk, t)
|
|
if add_to is not None:
|
|
res[:, t0:t1] += self.proj(out[:, t0:t1])
|
|
else:
|
|
res[:, t0:t1] = self.proj(out[:, t0:t1])
|
|
return res
|
|
|
|
|
|
class SwiGLU(nn.Module):
|
|
"""``w_down(silu(w_gate(x)) * w_up(x))``, chunked over tokens to bound the
|
|
``[chunk, hidden]`` workspace."""
|
|
|
|
def __init__(self, dim, hidden_dim):
|
|
super().__init__()
|
|
self.w_up = nn.Linear(dim, hidden_dim, bias=False)
|
|
self.w_gate = nn.Linear(dim, hidden_dim, bias=False)
|
|
self.w_down = nn.Linear(hidden_dim, dim, bias=False)
|
|
|
|
def forward(self, x, pre=None, add_to=None):
|
|
"""``pre``/``add_to`` as in ``NeighborhoodAttention3D.forward``."""
|
|
_, t, h, w, _ = x.shape
|
|
chunk = max(1, MLP_TOKEN_CHUNK // max(h * w, 1))
|
|
out = add_to if add_to is not None else torch.empty_like(x)
|
|
for t0 in range(0, t, chunk):
|
|
t1 = min(t0 + chunk, t)
|
|
sl = x[:, t0:t1] if pre is None else pre(x[:, t0:t1])
|
|
y = self.w_down(F.silu(self.w_gate(sl)) * self.w_up(sl))
|
|
if add_to is not None:
|
|
out[:, t0:t1] += y
|
|
else:
|
|
out[:, t0:t1] = y
|
|
return out
|
|
|
|
|
|
class NABlock(nn.Module):
|
|
"""Pre-norm transformer block: NA -> SwiGLU MLP with residual adds."""
|
|
|
|
def __init__(self, dim, kernel_size, head_dim=64, mlp_ratio=4.0):
|
|
super().__init__()
|
|
self.norm1 = RMSNorm(dim, eps=1e-6)
|
|
self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim)
|
|
self.norm2 = RMSNorm(dim, eps=1e-6)
|
|
hidden = (int(dim * mlp_ratio) + 15) // 16 * 16
|
|
self.mlp = SwiGLU(dim, hidden)
|
|
|
|
def forward(self, x):
|
|
x = self.attn(x, pre=self.norm1, add_to=x)
|
|
return self.mlp(x, pre=self.norm2, add_to=x)
|
|
|
|
|
|
def modulate(x, scale, shift):
|
|
return x * (1.0 + scale) + shift
|
|
|
|
|
|
class AdaLNZero(nn.Module):
|
|
"""``t_emb`` -> 7 (scale/shift/gate) chunks; gate slots unused (folded at export)."""
|
|
|
|
NUM_CHUNKS = 7
|
|
|
|
def __init__(self, dim, t_emb_dim):
|
|
super().__init__()
|
|
self.proj = nn.Linear(t_emb_dim, self.NUM_CHUNKS * dim, bias=True)
|
|
|
|
def forward(self, t_emb):
|
|
h = self.proj(F.silu(t_emb))
|
|
return tuple(c[:, None, None, None, :] for c in h.chunk(self.NUM_CHUNKS, dim=-1))
|
|
|
|
|
|
class DiffusionNABlock(nn.Module):
|
|
"""NA + SwiGLU with shared AdaLN-Zero scale/shift (ungated residuals)."""
|
|
|
|
def __init__(self, dim, kernel_size, context_channels, head_dim=64, mlp_ratio=4.0):
|
|
super().__init__()
|
|
self.context_proj = nn.Linear(context_channels, dim, bias=True)
|
|
self.scale_shift_table = nn.Parameter(torch.zeros(AdaLNZero.NUM_CHUNKS, dim))
|
|
self.norm1 = RMSNorm(dim, eps=1e-6)
|
|
self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim)
|
|
self.norm2 = RMSNorm(dim, eps=1e-6)
|
|
hidden = (int(dim * mlp_ratio) + 15) // 16 * 16
|
|
self.mlp = SwiGLU(dim, hidden)
|
|
|
|
def forward(self, x, latent_context, modulation):
|
|
scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [
|
|
modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(AdaLNZero.NUM_CHUNKS)
|
|
]
|
|
chunk = max(1, MLP_TOKEN_CHUNK // max(x.shape[2] * x.shape[3], 1))
|
|
for t0 in range(0, x.shape[1], chunk):
|
|
x[:, t0:t0 + chunk] += self.context_proj(latent_context[:, t0:t0 + chunk])
|
|
x = self.attn(x, pre=lambda s: modulate(self.norm1(s), scale_msa, shift_msa), add_to=x)
|
|
return self.mlp(x, pre=lambda s: modulate(self.norm2(s), scale_mlp, shift_mlp), add_to=x)
|
|
|
|
|
|
class LinearPixelShuffleUpsample(nn.Module):
|
|
"""Linear channel-expand, then channels-last pixel shuffle."""
|
|
|
|
def __init__(self, in_channels, stride, out_channels_reduction_factor=1):
|
|
super().__init__()
|
|
self.stride = tuple(stride)
|
|
proj_out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor
|
|
self.out_channels = proj_out_channels // math.prod(stride)
|
|
self.proj = nn.Linear(in_channels, proj_out_channels, bias=True)
|
|
|
|
def forward(self, x, drop_leading_frame=True):
|
|
batch, t, h, w, _ = x.shape
|
|
p1, p2, p3 = self.stride
|
|
out = torch.empty((batch, t * p1, h * p2, w * p3, self.out_channels), dtype=x.dtype, device=x.device)
|
|
chunk = max(1, MLP_TOKEN_CHUNK // max(h * w, 1))
|
|
for t0 in range(0, t, chunk):
|
|
t1 = min(t0 + chunk, t)
|
|
out[:, t0 * p1:t1 * p1] = rearrange(
|
|
self.proj(x[:, t0:t1]), "b t h w (c p1 p2 p3) -> b (t p1) (h p2) (w p3) c",
|
|
p1=p1, p2=p2, p3=p3,
|
|
)
|
|
if p1 == 2 and drop_leading_frame:
|
|
# The causal temporal pixel-shuffle duplicates the leading frame.
|
|
out = out[:, 1:]
|
|
return out
|
|
|
|
|
|
class TimestepEmbedder(nn.Module):
|
|
"""Sinusoidal(256) -> MLP. ``mlp.{0,2}`` naming matches the checkpoint."""
|
|
|
|
def __init__(self, t_emb_dim=384, freq_dim=256):
|
|
super().__init__()
|
|
self.freq_dim = freq_dim
|
|
self.mlp = nn.Sequential(
|
|
nn.Linear(freq_dim, t_emb_dim, bias=True),
|
|
nn.SiLU(),
|
|
nn.Linear(t_emb_dim, t_emb_dim, bias=True),
|
|
)
|
|
|
|
def forward(self, timestep, dtype):
|
|
emb = get_timestep_embedding(timestep.flatten(), self.freq_dim, flip_sin_to_cos=True,
|
|
downscale_freq_shift=0, scale=1)
|
|
return self.mlp(emb.to(dtype))
|
|
|
|
|
|
class NADiffusionDecoder(nn.Module):
|
|
"""Stages 1-4 (deterministic NA upsample) + stage-5 diffusion blocks.
|
|
|
|
Input latent must already be un-normalized (the wrapper applies
|
|
``per_channel_statistics.un_normalize``, same as the conv VAE path).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels=128,
|
|
out_channels=3,
|
|
patch_size=4,
|
|
head_dim=64,
|
|
stage_channels=(2048, 1024, 512, 512, 256),
|
|
stage_depths=(4, 6, 4, 2, 8),
|
|
stage_kernels=((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5), (11, 11, 11)),
|
|
upsamples=(((1, 2, 2), 2), ((2, 1, 1), 2), ((2, 2, 2), 1), ((2, 2, 2), 2)),
|
|
stage5_kernel=(11, 11, 11),
|
|
t_emb_dim=384,
|
|
default_num_inference_steps=1,
|
|
timestep_scale_multiplier=1000.0,
|
|
model_output_type="x0",
|
|
):
|
|
super().__init__()
|
|
self.patch_size = patch_size
|
|
self.out_channels = out_channels
|
|
self.timestep_scale_multiplier = timestep_scale_multiplier
|
|
self.model_output_type = model_output_type
|
|
self.register_buffer(
|
|
"default_inference_timesteps",
|
|
torch.linspace(1.0, 1.0 / default_num_inference_steps, default_num_inference_steps),
|
|
persistent=False,
|
|
)
|
|
self.temporal_upscale = math.prod(s[0] for s, _ in upsamples)
|
|
self.spatial_upscale = math.prod(s[1] for s, _ in upsamples) * patch_size
|
|
# NATTEN-style last-frame border mitigation: replicate the last latent
|
|
# frame through stages 1-4, crop the appendix off the context after.
|
|
self.trailing_pad_latent_frames = (stage_kernels[0][0] // 2) * 2
|
|
|
|
self.conv_in = nn.Linear(in_channels, stage_channels[0], bias=True)
|
|
|
|
self.det_stages = nn.ModuleList()
|
|
self.upsamples = nn.ModuleList()
|
|
for stage_i in range(len(stage_channels) - 1):
|
|
c = stage_channels[stage_i]
|
|
self.det_stages.append(nn.ModuleList(
|
|
[NABlock(c, stage_kernels[stage_i], head_dim=head_dim) for _ in range(stage_depths[stage_i])]
|
|
))
|
|
stride, reduction = upsamples[stage_i]
|
|
self.upsamples.append(LinearPixelShuffleUpsample(c, stride, out_channels_reduction_factor=reduction))
|
|
|
|
self.t_embedder = TimestepEmbedder(t_emb_dim=t_emb_dim)
|
|
|
|
c5 = stage_channels[-1]
|
|
self.context_channels = c5
|
|
noised_pixel_channels = out_channels * (patch_size ** 2)
|
|
self.conv_in_x_t = nn.Linear(noised_pixel_channels, c5, bias=True)
|
|
self.shared_adaln = AdaLNZero(c5, t_emb_dim)
|
|
self.diff_blocks = nn.ModuleList([
|
|
DiffusionNABlock(c5, stage5_kernel, context_channels=c5, head_dim=head_dim)
|
|
for _ in range(stage_depths[-1])
|
|
])
|
|
self.norm_out = RMSNorm(c5, eps=1e-6)
|
|
self.conv_out = nn.Linear(c5, noised_pixel_channels, bias=True)
|
|
|
|
def forward_pre_diffusion(self, z, drop_leading_frame=True, pad_trailing=True):
|
|
"""Stages 1-4: latent -> stage-5 context, channels-last.
|
|
|
|
``drop_leading_frame`` must be True only when ``z`` contains the
|
|
latent's true temporal origin (t=0); tiled callers decoding a later
|
|
temporal chunk pass False (the duplicate leading frame belongs solely
|
|
to the origin chunk). ``pad_trailing`` only for chunks containing the
|
|
latent's last frame."""
|
|
n = self.trailing_pad_latent_frames if pad_trailing else 0
|
|
if n > 0:
|
|
z = torch.cat([z, z[:, :, -1:].expand(-1, -1, n, -1, -1)], dim=2)
|
|
x = z.permute(0, 2, 3, 4, 1)
|
|
x = self.conv_in(x)
|
|
for stage_i, blocks in enumerate(self.det_stages):
|
|
for block in blocks:
|
|
x = block(x)
|
|
x = self.upsamples[stage_i](x, drop_leading_frame=drop_leading_frame)
|
|
if n > 0:
|
|
x = x[:, :-(n * self.temporal_upscale)]
|
|
return x
|
|
|
|
def forward_diff_step(self, context, x_t, t):
|
|
x = patchify(x_t, patch_size_hw=self.patch_size, patch_size_t=1)
|
|
x = self.conv_in_x_t(x.permute(0, 2, 3, 4, 1))
|
|
t_emb = self.t_embedder(self.timestep_scale_multiplier * t, dtype=x.dtype)
|
|
modulation = self.shared_adaln(t_emb)
|
|
for block in self.diff_blocks:
|
|
x = block(x, context, modulation)
|
|
x = self.norm_out(x)
|
|
x = self.conv_out(x)
|
|
x = x.permute(0, 4, 1, 2, 3)
|
|
return unpatchify(x, patch_size_hw=self.patch_size, patch_size_t=1)
|
|
|
|
def forward(self, z, generator=None, drop_leading_frame=True, pad_trailing=True):
|
|
context = self.forward_pre_diffusion(z, drop_leading_frame=drop_leading_frame, pad_trailing=pad_trailing)
|
|
batch, t5, h5, w5, _ = context.shape
|
|
pixel_shape = (batch, self.out_channels, t5, h5 * self.patch_size, w5 * self.patch_size)
|
|
x_t = torch.randn(pixel_shape, dtype=z.dtype, device=z.device, generator=generator)
|
|
|
|
timesteps = self.default_inference_timesteps.to(z.device)
|
|
num_steps = timesteps.shape[0]
|
|
for i in range(num_steps):
|
|
t_now = timesteps[i].expand(batch)
|
|
model_out = self.forward_diff_step(context, x_t, t_now)
|
|
if self.model_output_type == "x0":
|
|
x0 = model_out
|
|
if i == num_steps - 1:
|
|
return x0
|
|
velocity = (x_t.float() - x0.float()) / timesteps[i]
|
|
else: # "v"
|
|
velocity = model_out.float()
|
|
if i == num_steps - 1:
|
|
return (x_t.float() - timesteps[i] * velocity).to(z.dtype)
|
|
t_next = timesteps[i + 1] if i + 1 < num_steps else torch.zeros_like(timesteps[i])
|
|
x_t = (x_t.float() - (timesteps[i] - t_next) * velocity).to(z.dtype)
|
|
return x_t
|
|
|
|
|
|
LTX_24_VAE_CONFIG = {
|
|
"_class_name": "CausalDiffusionVAE",
|
|
"dims": 3,
|
|
"model_output_type": "x0",
|
|
"encoder": {
|
|
"dims": 3,
|
|
"in_channels": 3,
|
|
"out_channels": 128,
|
|
"blocks": [
|
|
["res_x", {"num_layers": 4}],
|
|
["compress_space_res", {"multiplier": 2}],
|
|
["res_x", {"num_layers": 6}],
|
|
["compress_time_res", {"multiplier": 2}],
|
|
["res_x", {"num_layers": 4}],
|
|
["compress_all_res", {"multiplier": 2}],
|
|
["res_x", {"num_layers": 2}],
|
|
["compress_all_res", {"multiplier": 1}],
|
|
["res_x", {"num_layers": 2}],
|
|
],
|
|
"patch_size": 4,
|
|
"latent_log_var": "constant",
|
|
"norm_layer": "pixel_norm",
|
|
"base_channels": 128,
|
|
"spatial_padding_mode": "zeros",
|
|
},
|
|
"decoder": {
|
|
"in_channels": 128,
|
|
"out_channels": 3,
|
|
"patch_size": 4,
|
|
"head_dim": 64,
|
|
"stage_channels": [2048, 1024, 512, 512, 256],
|
|
"stage_depths": [4, 6, 4, 2, 8],
|
|
"stage_kernels": [[3, 7, 7], [3, 7, 7], [3, 5, 5], [3, 5, 5], [11, 11, 11]],
|
|
"upsamples": [[[1, 2, 2], 2], [[2, 1, 1], 2], [[2, 2, 2], 1], [[2, 2, 2], 2]],
|
|
"stage5_kernel": [11, 11, 11],
|
|
"timestep_scale_multiplier": 1000.0,
|
|
"default_num_inference_steps": 1,
|
|
},
|
|
}
|
|
|
|
|
|
class CausalDiffusionVAE(nn.Module):
|
|
"""LTX 2.4 video VAE: conv encoder (shared with the 2.0 arch) + NA
|
|
diffusion decoder. Interface mirrors ``causal_video_autoencoder.VideoVAE``.
|
|
"""
|
|
|
|
def __init__(self, config=None):
|
|
super().__init__()
|
|
if config is None:
|
|
config = LTX_24_VAE_CONFIG
|
|
self.config = config
|
|
enc = config.get("encoder", LTX_24_VAE_CONFIG["encoder"])
|
|
dec = config.get("decoder", LTX_24_VAE_CONFIG["decoder"])
|
|
dec_defaults = LTX_24_VAE_CONFIG["decoder"]
|
|
|
|
self.encoder = Encoder(
|
|
dims=enc.get("dims", 3),
|
|
in_channels=enc.get("in_channels", 3),
|
|
out_channels=enc.get("out_channels", 128),
|
|
blocks=enc.get("blocks", LTX_24_VAE_CONFIG["encoder"]["blocks"]),
|
|
patch_size=enc.get("patch_size", 4),
|
|
latent_log_var=enc.get("latent_log_var", "constant"),
|
|
norm_layer=enc.get("norm_layer", "pixel_norm"),
|
|
spatial_padding_mode=enc.get("spatial_padding_mode", "zeros"),
|
|
base_channels=enc.get("base_channels", 128),
|
|
)
|
|
|
|
self.decoder = NADiffusionDecoder(
|
|
in_channels=dec.get("in_channels", 128),
|
|
out_channels=dec.get("out_channels", 3),
|
|
patch_size=dec.get("patch_size", 4),
|
|
head_dim=dec.get("head_dim", 64),
|
|
stage_channels=tuple(dec.get("stage_channels", dec_defaults["stage_channels"])),
|
|
stage_depths=tuple(dec.get("stage_depths", dec_defaults["stage_depths"])),
|
|
stage_kernels=tuple(tuple(k) for k in dec.get("stage_kernels", dec_defaults["stage_kernels"])),
|
|
upsamples=tuple((tuple(s), r) for s, r in dec.get("upsamples", dec_defaults["upsamples"])),
|
|
stage5_kernel=tuple(dec.get("stage5_kernel", dec_defaults["stage5_kernel"])),
|
|
t_emb_dim=dec.get("t_emb_dim", 384),
|
|
default_num_inference_steps=dec.get("default_num_inference_steps", 1),
|
|
timestep_scale_multiplier=dec.get("timestep_scale_multiplier", 1000.0),
|
|
model_output_type=config.get("model_output_type", "x0"),
|
|
)
|
|
|
|
self.per_channel_statistics = processor()
|
|
|
|
def encode(self, x, device=None):
|
|
x = x[:, :, :max(1, 1 + ((x.shape[2] - 1) // 8) * 8), :, :]
|
|
means, logvar = torch.chunk(self.encoder(x, device=device), 2, dim=1)
|
|
return self.per_channel_statistics.normalize(means)
|
|
|
|
def decode(self, x):
|
|
# Fixed-seed noise so decodes are reproducible TODO: expose?
|
|
generator = torch.Generator(device=x.device)
|
|
generator.manual_seed(0)
|
|
return self.decoder(self.per_channel_statistics.un_normalize(x), generator=generator)
|