Skip to content

model_autoencoders_shared model/pipeline review #13652

Open
Open
model_autoencoders_shared model/pipeline review#13652

Description

@hlky
hlky
opened on Apr 29, 2026

model_autoencoders_shared model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Full target file list reviewed: autoencoder_asym_kl.py, autoencoder_dc.py, autoencoder_kl.py, autoencoder_kl_kvae.py, autoencoder_kl_kvae_video.py, autoencoder_kl_magvit.py, autoencoder_kl_temporal_decoder.py, autoencoder_oobleck.py, autoencoder_rae.py, autoencoder_tiny.py, autoencoder_vidtok.py, consistency_decoder_vae.py, vae.py, vq_model.py.

Issue 1: AutoencoderKLKVAEVideo discards encoder log variance

Affected code:

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_kvae_video.py

Lines 848 to 894 in 0f1abc4

def _encode(self, x: torch.Tensor, seg_len: int = 16) -> torch.Tensor:
# Cached encoder processes by segments
cache = self._make_encoder_cache()
split_list = [seg_len + 1]
n_frames = x.size(2) - (seg_len + 1)
while n_frames > 0:
split_list.append(seg_len)
n_frames -= seg_len
split_list[-1] += n_frames
latent = []
for chunk in torch.split(x, split_list, dim=2):
l = self.encoder(chunk, cache)
sample, _ = torch.chunk(l, 2, dim=1)
latent.append(sample)
return torch.cat(latent, dim=2)
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
"""
Encode a batch of videos into latents.
Args:
x (`torch.Tensor`): Input batch of videos with shape (B, C, T, H, W).
return_dict (`bool`, *optional*, defaults to `True`):
Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
Returns:
The latent representations of the encoded videos.
"""
if self.use_slicing and x.shape[0] > 1:
encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
h = torch.cat(encoded_slices)
else:
h = self._encode(x)
# For cached encoder, we already did the split in _encode
h_double = torch.cat([h, torch.zeros_like(h)], dim=1)
posterior = DiagonalGaussianDistribution(h_double)
if not return_dict:
return (posterior,)
return AutoencoderKLOutput(latent_dist=posterior)

Problem:
KVAECachedEncoder3D outputs 2 * z_channels, but _encode() splits the tensor and keeps only the first half. encode() then reconstructs a fake [mean, zeros] tensor, so the posterior log variance is always zero.

Impact:
sample_posterior=True samples from the wrong distribution, checkpoint log-variance weights are ignored, and parity with the source KVAE 3D VAE is broken.

Reproduction:

import torch
from diffusers import AutoencoderKLKVAEVideo

model = AutoencoderKLKVAEVideo(
    ch=32, ch_mult=(1, 2), num_res_blocks=1, z_channels=4, temporal_compress_times=2
).eval()
x = torch.randn(1, 3, 3, 16, 16)

with torch.no_grad():
    raw = model.encoder(x, model._make_encoder_cache())
    raw_mean, raw_logvar = raw.chunk(2, dim=1)
    posterior = model.encode(x).latent_dist

print(torch.allclose(posterior.mean, raw_mean, atol=1e-5))
print(raw_logvar.abs().max().item())
print(posterior.logvar.abs().max().item())  # always 0.0

Relevant precedent:

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_kvae.py

Lines 592 to 629 in 0f1abc4

def _encode(self, x: torch.Tensor) -> torch.Tensor:
batch_size, num_channels, height, width = x.shape
if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size):
return self._tiled_encode(x)
enc = self.encoder(x)
return enc
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
"""
Encode a batch of images into latents.
Args:
x (`torch.Tensor`): Input batch of images.
return_dict (`bool`, *optional*, defaults to `True`):
Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
Returns:
The latent representations of the encoded images. If `return_dict` is True, a
[`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
"""
if self.use_slicing and x.shape[0] > 1:
encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
h = torch.cat(encoded_slices)
else:
h = self._encode(x)
posterior = DiagonalGaussianDistribution(h)
if not return_dict:
return (posterior,)
return AutoencoderKLOutput(latent_dist=posterior)

Suggested fix:

# _encode(): keep the full encoder output
latent.append(self.encoder(chunk, cache))

# encode(): do not synthesize a zero logvar half
posterior = DiagonalGaussianDistribution(h)

Issue 2: AutoencoderVidTok(return_dict=False) returns a tensor instead of a tuple

Affected code:

diffusers/src/diffusers/models/autoencoders/autoencoder_vidtok.py

Lines 1435 to 1488 in 0f1abc4

def forward(
self,
sample: torch.Tensor,
sample_posterior: bool = True,
encoder_mode: bool = False,
return_dict: bool = True,
generator: Optional[torch.Generator] = None,
) -> Union[torch.Tensor, DecoderOutput]:
x = sample
res = 1 if self.is_causal else 0
if self.is_causal:
if x.shape[2] % self.temporal_compression_ratio != res:
time_padding = self.temporal_compression_ratio - x.shape[2] % self.temporal_compression_ratio + res
x = self._pad_at_dim(x, (0, time_padding), dim=2, pad_mode="replicate")
else:
time_padding = 0
else:
if x.shape[2] % self.num_sample_frames_batch_size != res:
if not encoder_mode:
time_padding = (
self.num_sample_frames_batch_size - x.shape[2] % self.num_sample_frames_batch_size + res
)
x = self._pad_at_dim(x, (0, time_padding), dim=2, pad_mode="replicate")
else:
assert x.shape[2] >= self.num_sample_frames_batch_size, (
f"Too short video. At least {self.num_sample_frames_batch_size} frames."
)
x = x[:, :, : x.shape[2] // self.num_sample_frames_batch_size * self.num_sample_frames_batch_size]
else:
time_padding = 0
if self.is_causal:
x = self._pad_at_dim(x, (self.temporal_compression_ratio - 1, 0), dim=2, pad_mode="replicate")
if self.regularizer == "kl":
posterior = self.encode(x).latent_dist
if sample_posterior:
z = posterior.sample(generator=generator)
else:
z = posterior.mode()
if encoder_mode:
return z
else:
z, indices = self.encode(x)
if encoder_mode:
return z, indices
dec = self.decode(z)
if time_padding != 0:
dec = dec[:, :, :-time_padding, :, :]
if not return_dict:
return dec
return DecoderOutput(sample=dec)

Problem:
Every other autoencoder forward path returns a tuple when return_dict=False. VidTok returns the raw decoded tensor, so common caller code like model(..., return_dict=False)[0] silently selects the first batch element.

Impact:
Pipeline/model utility code that relies on diffusers’ tuple convention gets the wrong shape and silently drops batch items.

Reproduction:

import torch
from diffusers import AutoencoderVidTok

model = AutoencoderVidTok(
    is_causal=False, ch=8, ch_mult=[1], z_channels=2, double_z=True, num_res_blocks=1, regularizer="kl"
).eval()
x = torch.randn(2, 3, 2, 8, 8)

with torch.no_grad():
    out = model(x, sample_posterior=False, return_dict=False)

print(type(out), out.shape)      # torch.Tensor, not tuple
print(out[0].shape)              # first batch item, not decoded tuple element

Relevant precedent:

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_kvae_video.py

Lines 938 to 954 in 0f1abc4

def forward(
self,
sample: torch.Tensor,
sample_posterior: bool = False,
return_dict: bool = True,
generator: Optional[torch.Generator] = None,
) -> Union[DecoderOutput, torch.Tensor]:
x = sample
posterior = self.encode(x).latent_dist
if sample_posterior:
z = posterior.sample(generator=generator)
else:
z = posterior.mode()
dec = self.decode(z).sample
if not return_dict:
return (dec,)
return DecoderOutput(sample=dec)

Suggested fix:

if not return_dict:
    return (dec,)
return DecoderOutput(sample=dec)

Issue 3: AutoencoderVidTok bypasses diffusers attention processors

Affected code:

diffusers/src/diffusers/models/autoencoders/autoencoder_vidtok.py

Lines 426 to 447 in 0f1abc4

class VidTokAttnBlock(nn.Module):
r"""A 3D self-attention block used in VidTok Model."""
def __init__(self, in_channels: int, is_causal: bool = True):
super().__init__()
make_conv_cls = VidTokCausalConv3d if is_causal else nn.Conv3d
self.norm = VidTokLayerNorm(dim=in_channels, eps=1e-6)
self.q = make_conv_cls(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
self.k = make_conv_cls(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
self.v = make_conv_cls(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
self.proj_out = make_conv_cls(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
def attention(self, hidden_states: torch.Tensor) -> torch.Tensor:
r"""Implement self-attention."""
hidden_states = self.norm(hidden_states)
q = self.q(hidden_states)
k = self.k(hidden_states)
v = self.v(hidden_states)
b, c, t, h, w = q.shape
q, k, v = [x.permute(0, 2, 3, 4, 1).reshape(b, t, -1, c).contiguous() for x in [q, k, v]]
hidden_states = F.scaled_dot_product_attention(q, k, v) # scale is dim ** -0.5 per default
return hidden_states.reshape(b, t, h, w, c).permute(0, 4, 1, 2, 3)

diffusers/src/diffusers/models/autoencoders/autoencoder_vidtok.py

Lines 938 to 979 in 0f1abc4

class AutoencoderVidTok(ModelMixin, ConfigMixin):
r"""
A VAE model for encoding videos into latents and decoding latent representations into videos, supporting both
continuous and discrete latent representations. Used in [VidTok](https://github.com/microsoft/VidTok).
This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
for all models (such as downloading or saving).
Args:
in_channels (`int`, defaults to 3):
The number of input channels.
out_channels (`int`, defaults to 3):
The number of output channels.
ch (`int`, defaults to 128):
The number of the basic channel.
ch_mult (`List[int]`, defaults to `[1, 2, 4, 4]`):
The multiple of the basic channel for each block.
z_channels (`int`, defaults to 4):
The number of latent channels.
double_z (`bool`, defaults to `True`):
Whether or not to double the z_channels.
num_res_blocks (`int`, defaults to 2):
The number of resblocks.
spatial_ds (`List`, *optional*, defaults to `None`):
Spatial downsample layers.
spatial_us (`List`, *optional*, defaults to `None`):
Spatial upsample layers.
tempo_ds (`List`, *optional*, defaults to `None`):
Temporal downsample layers.
tempo_us (`List`, *optional*, defaults to `None`):
Temporal upsample layers.
dropout (`float`, defaults to 0.0):
Dropout rate.
regularizer (`str`, defaults to `"kl"`):
The regularizer type - "kl" for continuous cases and "fsq" for discrete cases.
codebook_size (`int`, defaults to 262144):
The codebook size used only in discrete cases.
is_causal (`bool`, defaults to `True`):
Whether it is a causal module.
"""
_supports_gradient_checkpointing = True

Problem:
VidTokAttnBlock calls F.scaled_dot_product_attention directly, and AutoencoderVidTok does not inherit AttentionMixin. This bypasses the repository’s attention processor/backend path.

Impact:
Users cannot inspect or replace attention processors, attention backend selection is ignored for VidTok attention, and the implementation violates the model attention rule.

Reproduction:

from diffusers import AutoencoderVidTok

model = AutoencoderVidTok(
    is_causal=False, ch=8, ch_mult=[1], z_channels=2, double_z=True, num_res_blocks=1, regularizer="kl"
)

print(hasattr(model, "set_attn_processor"))  # False
print([type(m).__name__ for m in model.modules() if type(m).__name__ == "VidTokAttnBlock"])

Relevant precedent:

diffusers/src/diffusers/models/autoencoders/autoencoder_rae.py

Lines 202 to 208 in 0f1abc4

self.attention = Attention(
query_dim=hidden_size,
heads=num_attention_heads,
dim_head=hidden_size // num_attention_heads,
dropout=attention_probs_dropout_prob,
bias=qkv_bias,
)

diffusers/src/diffusers/models/autoencoders/autoencoder_rae.py

Lines 393 to 442 in 0f1abc4

class AutoencoderRAE(ModelMixin, AttentionMixin, AutoencoderMixin, ConfigMixin):
r"""
Representation Autoencoder (RAE) model for encoding images to latents and decoding latents to images.
This model uses a frozen pretrained encoder (DINOv2, SigLIP2, or MAE) with a trainable ViT decoder to reconstruct
images from learned representations.
This model inherits from [`ModelMixin`]. Check the superclass documentation for its generic methods implemented for
all models (such as downloading or saving).
Args:
encoder_type (`str`, *optional*, defaults to `"dinov2"`):
Type of frozen encoder to use. One of `"dinov2"`, `"siglip2"`, or `"mae"`.
encoder_hidden_size (`int`, *optional*, defaults to `768`):
Hidden size of the encoder model.
encoder_patch_size (`int`, *optional*, defaults to `14`):
Patch size of the encoder model.
encoder_num_hidden_layers (`int`, *optional*, defaults to `12`):
Number of hidden layers in the encoder model.
patch_size (`int`, *optional*, defaults to `16`):
Decoder patch size (used for unpatchify and decoder head).
encoder_input_size (`int`, *optional*, defaults to `224`):
Input size expected by the encoder.
image_size (`int`, *optional*):
Decoder output image size. If `None`, it is derived from encoder token count and `patch_size` like
RAE-main: `image_size = patch_size * sqrt(num_patches)`, where `num_patches = (encoder_input_size //
encoder_patch_size) ** 2`.
num_channels (`int`, *optional*, defaults to `3`):
Number of input/output channels.
encoder_norm_mean (`list`, *optional*, defaults to `[0.485, 0.456, 0.406]`):
Channel-wise mean for encoder input normalization (ImageNet defaults).
encoder_norm_std (`list`, *optional*, defaults to `[0.229, 0.224, 0.225]`):
Channel-wise std for encoder input normalization (ImageNet defaults).
latents_mean (`list` or `tuple`, *optional*):
Optional mean for latent normalization. Tensor inputs are accepted and converted to config-serializable
lists.
latents_std (`list` or `tuple`, *optional*):
Optional standard deviation for latent normalization. Tensor inputs are accepted and converted to
config-serializable lists.
noise_tau (`float`, *optional*, defaults to `0.0`):
Noise level for training (adds noise to latents during training).
reshape_to_2d (`bool`, *optional*, defaults to `True`):
Whether to reshape latents to 2D (B, C, H, W) format.
use_encoder_loss (`bool`, *optional*, defaults to `False`):
Whether to use encoder hidden states in the loss (for advanced training).
"""
# NOTE: gradient checkpointing is not wired up for this model yet.
_supports_gradient_checkpointing = False
_no_split_modules = ["ViTMAELayer"]

Suggested fix:
Refactor VidTokAttnBlock to use the diffusers attention processor pattern with dispatch_attention_fn, and make AutoencoderVidTok inherit AttentionMixin.

Issue 4: AutoencoderTiny fails bfloat16 forward because it creates float32 activations

Affected code:

diffusers/src/diffusers/models/autoencoders/autoencoder_tiny.py

Lines 302 to 312 in 0f1abc4

enc = self.encode(sample).latents
# scale latents to be in [0, 1], then quantize latents to a byte tensor,
# as if we were storing the latents in an RGBA uint8 image.
scaled_enc = self.scale_latents(enc).mul_(255).round_().byte()
# unquantize latents back into [0, 1], then unscale latents back to their original range,
# as if we were loading the latents from an RGBA uint8 image.
unscaled_enc = self.unscale_latents(scaled_enc / 255.0)
dec = self.decode(unscaled_enc).sample

diffusers/tests/models/autoencoders/test_models_autoencoder_tiny.py

Lines 146 to 159 in 0f1abc4

@unittest.skip(
"The forward pass of AutoencoderTiny creates a torch.float32 tensor. This causes inference in compute_dtype=torch.bfloat16 to fail. To fix:\n"
"1. Change the forward pass to be dtype agnostic.\n"
"2. Unskip this test."
)
def test_layerwise_casting_inference(self):
pass
@unittest.skip(
"The forward pass of AutoencoderTiny creates a torch.float32 tensor. This causes inference in compute_dtype=torch.bfloat16 to fail. To fix:\n"
"1. Change the forward pass to be dtype agnostic.\n"
"2. Unskip this test."
)
def test_layerwise_casting_memory(self):

Problem:
scaled_enc / 255.0 promotes the byte tensor to float32 before decode. With bfloat16 weights, the decoder receives float32 inputs and errors. The fast tests already skip layerwise-casting coverage with this exact reason.

Impact:
Layerwise casting / bfloat16 inference cannot cover AutoencoderTiny, and users hit dtype mismatch errors.

Reproduction:

import torch
from diffusers import AutoencoderTiny

model = AutoencoderTiny(
    encoder_block_out_channels=(8, 8, 8, 8),
    decoder_block_out_channels=(8, 8, 8, 8),
    num_encoder_blocks=(1, 1, 1, 1),
    num_decoder_blocks=(1, 1, 1, 1),
).eval().to(dtype=torch.bfloat16)

x = torch.randn(1, 3, 32, 32, dtype=torch.bfloat16)
model(x)

Relevant precedent:
The skipped tests at the link above already describe the expected fix.

Suggested fix:

unscaled_enc = self.unscale_latents(scaled_enc.to(dtype=enc.dtype) / 255.0)

Issue 5: AutoencoderTiny tiled paths hardcode channel counts and derive scale from out_channels

Affected code:

diffusers/src/diffusers/models/autoencoders/autoencoder_tiny.py

Lines 147 to 151 in 0f1abc4

# only relevant if vae tiling is enabled
self.spatial_scale_factor = 2**out_channels
self.tile_overlap_factor = 0.125
self.tile_sample_min_size = 512
self.tile_latent_min_size = self.tile_sample_min_size // self.spatial_scale_factor

diffusers/src/diffusers/models/autoencoders/autoencoder_tiny.py

Lines 190 to 196 in 0f1abc4

blend_masks = torch.stack(
torch.meshgrid([torch.arange(tile_size / sf) / (blend_size / sf - 1)] * 2, indexing="ij")
)
blend_masks = blend_masks.clamp(0, 1).to(x.device)
# output array
out = torch.zeros(x.shape[0], 4, x.shape[-2] // sf, x.shape[-1] // sf, device=x.device)

diffusers/src/diffusers/models/autoencoders/autoencoder_tiny.py

Lines 238 to 244 in 0f1abc4

blend_masks = torch.stack(
torch.meshgrid([torch.arange(tile_size * sf) / (blend_size * sf - 1)] * 2, indexing="ij")
)
blend_masks = blend_masks.clamp(0, 1).to(x.device)
# output array
out = torch.zeros(x.shape[0], 3, x.shape[-2] * sf, x.shape[-1] * sf, device=x.device)

Problem:
Tiling assumes latent channels are always 4, output channels are always 3, and spatial scale is 2 ** out_channels. These are configurable constructor values.

Impact:
Non-default but valid AutoencoderTiny configs fail as soon as tiling is enabled.

Reproduction:

import torch
from diffusers import AutoencoderTiny

model = AutoencoderTiny(
    out_channels=1,
    encoder_block_out_channels=(8, 8, 8, 8),
    decoder_block_out_channels=(8, 8, 8, 8),
    num_encoder_blocks=(1, 1, 1, 1),
    num_decoder_blocks=(1, 1, 1, 1),
    latent_channels=4,
)
model.enable_tiling()

x = torch.randn(1, 3, 32, 32)
model.encode(x)

Relevant precedent:
Other tiled autoencoders derive shapes from tensor/model config rather than fixed RGB/latent-channel constants.

Suggested fix:

self.spatial_scale_factor = 2 ** (len(encoder_block_out_channels) - 1)

out = torch.zeros(
    x.shape[0], self.config.latent_channels,
    x.shape[-2] // sf, x.shape[-1] // sf,
    device=x.device, dtype=x.dtype,
)

out = torch.zeros(
    x.shape[0], self.config.out_channels,
    x.shape[-2] * sf, x.shape[-1] * sf,
    device=x.device, dtype=x.dtype,
)

Issue 6: Tuple sample_size breaks VAE tiling comparisons

Affected code:

diffusers/src/diffusers/models/autoencoders/autoencoder_kl.py

Lines 132 to 162 in 0f1abc4

# only relevant if vae tiling is enabled
self.tile_sample_min_size = self.config.sample_size
sample_size = (
self.config.sample_size[0]
if isinstance(self.config.sample_size, (list, tuple))
else self.config.sample_size
)
self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.block_out_channels) - 1)))
self.tile_overlap_factor = 0.25
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
def set_default_attn_processor(self):
"""
Disables custom attention processors and sets the default attention implementation.
"""
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
processor = AttnAddedKVProcessor()
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
processor = AttnProcessor()
else:
raise ValueError(
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
)
self.set_attn_processor(processor)
def _encode(self, x: torch.Tensor) -> torch.Tensor:
batch_size, num_channels, height, width = x.shape
if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size):
return self._tiled_encode(x)

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_kvae.py

Lines 582 to 596 in 0f1abc4

# only relevant if vae tiling is enabled
self.tile_sample_min_size = self.config.sample_size
sample_size = (
self.config.sample_size[0]
if isinstance(self.config.sample_size, (list, tuple))
else self.config.sample_size
)
self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.ch_mult) - 1)))
self.tile_overlap_factor = 0.25
def _encode(self, x: torch.Tensor) -> torch.Tensor:
batch_size, num_channels, height, width = x.shape
if self.use_tiling and (width > self.tile_sample_min_size or height > self.tile_sample_min_size):
return self._tiled_encode(x)

diffusers/src/diffusers/models/autoencoders/consistency_decoder_vae.py

Lines 159 to 204 in 0f1abc4

# only relevant if vae tiling is enabled
self.tile_sample_min_size = self.config.sample_size
sample_size = (
self.config.sample_size[0]
if isinstance(self.config.sample_size, (list, tuple))
else self.config.sample_size
)
self.tile_latent_min_size = int(sample_size / (2 ** (len(self.config.block_out_channels) - 1)))
self.tile_overlap_factor = 0.25
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
def set_default_attn_processor(self):
"""
Disables custom attention processors and sets the default attention implementation.
"""
if all(proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
processor = AttnAddedKVProcessor()
elif all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
processor = AttnProcessor()
else:
raise ValueError(
f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
)
self.set_attn_processor(processor)
@apply_forward_hook
def encode(
self, x: torch.Tensor, return_dict: bool = True
) -> ConsistencyDecoderVAEOutput | tuple[DiagonalGaussianDistribution]:
"""
Encode a batch of images into latents.
Args:
x (`torch.Tensor`): Input batch of images.
return_dict (`bool`, *optional*, defaults to `True`):
Whether to return a [`~models.autoencoders.consistency_decoder_vae.ConsistencyDecoderVAEOutput`]
instead of a plain tuple.
Returns:
The latent representations of the encoded images. If `return_dict` is True, a
[`~models.autoencoders.consistency_decoder_vae.ConsistencyDecoderVAEOutput`] is returned, otherwise a
plain `tuple` is returned.
"""
if self.use_tiling and (x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size):
return self.tiled_encode(x, return_dict=return_dict)

Problem:
The constructors partly support sample_size as a tuple/list when computing latent tile size, but leave self.tile_sample_min_size as the tuple. Later tiling checks compare int > tuple.

Impact:
Loading or constructing rectangular VAE configs with tuple sample sizes crashes when tiling is enabled.

Reproduction:

import torch
from diffusers import AutoencoderKL

model = AutoencoderKL(
    sample_size=(32, 64),
    block_out_channels=(4,),
    norm_num_groups=1,
    latent_channels=2,
)
model.enable_tiling()
model.encode(torch.randn(1, 3, 65, 65))

Relevant precedent:
AutoencoderDC keeps separate height/width tile thresholds.

Suggested fix:

sample_size = self.config.sample_size
if isinstance(sample_size, (list, tuple)):
    self.tile_sample_min_height, self.tile_sample_min_width = sample_size
else:
    self.tile_sample_min_height = self.tile_sample_min_width = sample_size

if self.use_tiling and (width > self.tile_sample_min_width or height > self.tile_sample_min_height):
    ...

Issue 7: Several target models have no slow model-level tests

Affected code:

diffusers/tests/models/autoencoders/test_models_autoencoder_dc.py

Lines 81 to 104 in 0f1abc4

class TestAutoencoderDC(AutoencoderDCTesterConfig, ModelTesterMixin):
base_precision = 1e-2
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype):
if dtype == torch.bfloat16 and IS_GITHUB_ACTIONS:
pytest.skip("Skipping bf16 test inside GitHub Actions environment")
super().test_from_save_pretrained_dtype_inference(tmp_path, dtype)
class TestAutoencoderDCTraining(AutoencoderDCTesterConfig, TrainingTesterMixin):
"""Training tests for AutoencoderDC."""
class TestAutoencoderDCMemory(AutoencoderDCTesterConfig, MemoryTesterMixin):
"""Memory optimization tests for AutoencoderDC."""
@pytest.mark.skipif(IS_GITHUB_ACTIONS, reason="Skipping test inside GitHub Actions environment")
def test_layerwise_casting_memory(self):
super().test_layerwise_casting_memory()
class TestAutoencoderDCSlicingTiling(AutoencoderDCTesterConfig, NewAutoencoderTesterMixin):
"""Slicing and tiling tests for AutoencoderDC."""

diffusers/tests/models/autoencoders/test_models_autoencoder_kl_kvae.py

Lines 28 to 73 in 0f1abc4

class AutoencoderKLKVAETests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase):
model_class = AutoencoderKLKVAE
main_input_name = "sample"
base_precision = 1e-2
def get_autoencoder_kl_kvae_config(self):
return {
"in_channels": 3,
"channels": 32,
"num_enc_blocks": 1,
"num_dec_blocks": 1,
"z_channels": 4,
"double_z": True,
"ch_mult": (1, 2),
"sample_size": 32,
}
@property
def dummy_input(self):
batch_size = 2
num_channels = 3
sizes = (32, 32)
image = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
return {"sample": image}
@property
def input_shape(self):
return (3, 32, 32)
@property
def output_shape(self):
return (3, 32, 32)
def prepare_init_args_and_inputs_for_common(self):
init_dict = self.get_autoencoder_kl_kvae_config()
inputs_dict = self.dummy_input
return init_dict, inputs_dict
def test_gradient_checkpointing_is_applied(self):
expected_set = {
"KVAEEncoder2D",
"KVAEDecoder2D",
}
super().test_gradient_checkpointing_is_applied(expected_set=expected_set)

diffusers/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py

Lines 28 to 118 in 0f1abc4

class AutoencoderKLKVAEVideoTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase):
model_class = AutoencoderKLKVAEVideo
main_input_name = "sample"
base_precision = 1e-2
def get_autoencoder_kl_kvae_video_config(self):
return {
"ch": 32,
"ch_mult": (1, 2),
"num_res_blocks": 1,
"in_channels": 3,
"out_ch": 3,
"z_channels": 4,
"temporal_compress_times": 2,
}
@property
def dummy_input(self):
batch_size = 2
num_frames = 3 # satisfies (T-1) % temporal_compress_times == 0 with temporal_compress_times=2
num_channels = 3
sizes = (16, 16)
video = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device)
return {"sample": video}
@property
def input_shape(self):
return (3, 3, 16, 16)
@property
def output_shape(self):
return (3, 3, 16, 16)
def prepare_init_args_and_inputs_for_common(self):
init_dict = self.get_autoencoder_kl_kvae_video_config()
inputs_dict = self.dummy_input
return init_dict, inputs_dict
def test_gradient_checkpointing_is_applied(self):
expected_set = {
"KVAECachedEncoder3D",
"KVAECachedDecoder3D",
}
super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
@unittest.skip("Unsupported test.")
def test_outputs_equivalence(self):
pass
@unittest.skip(
"Multi-GPU inference is not supported due to the stateful cache_dict passing through the forward pass."
)
def test_model_parallelism(self):
pass
@unittest.skip(
"Multi-GPU inference is not supported due to the stateful cache_dict passing through the forward pass."
)
def test_sharded_checkpoints_device_map(self):
pass
def _run_nondeterministic(self, fn):
# reflection_pad3d_backward_out_cuda has no deterministic CUDA implementation;
# temporarily relax the requirement for training tests that do backward passes.
import torch
torch.use_deterministic_algorithms(False)
try:
fn()
finally:
torch.use_deterministic_algorithms(True)
def test_training(self):
self._run_nondeterministic(super().test_training)
def test_ema_training(self):
self._run_nondeterministic(super().test_ema_training)
@unittest.skip(
"Gradient checkpointing recomputes the forward pass, but the model uses a stateful cache_dict "
"that is mutated during the first forward. On recomputation the cache is already populated, "
"causing a different execution path and numerically different gradients. "
"GC still reduces peak memory usage; gradient correctness in the presence of GC is a known limitation."
)
def test_effective_gradient_checkpointing(self):
pass
def test_layerwise_casting_training(self):
self._run_nondeterministic(super().test_layerwise_casting_training)

diffusers/tests/models/autoencoders/test_models_autoencoder_magvit.py

Lines 28 to 97 in 0f1abc4

class AutoencoderKLMagvitTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase):
model_class = AutoencoderKLMagvit
main_input_name = "sample"
base_precision = 1e-2
def get_autoencoder_kl_magvit_config(self):
return {
"in_channels": 3,
"latent_channels": 4,
"out_channels": 3,
"block_out_channels": [8, 8, 8, 8],
"down_block_types": [
"SpatialDownBlock3D",
"SpatialTemporalDownBlock3D",
"SpatialTemporalDownBlock3D",
"SpatialTemporalDownBlock3D",
],
"up_block_types": [
"SpatialUpBlock3D",
"SpatialTemporalUpBlock3D",
"SpatialTemporalUpBlock3D",
"SpatialTemporalUpBlock3D",
],
"layers_per_block": 1,
"norm_num_groups": 8,
"spatial_group_norm": True,
}
@property
def dummy_input(self):
batch_size = 2
num_frames = 9
num_channels = 3
height = 16
width = 16
image = floats_tensor((batch_size, num_channels, num_frames, height, width)).to(torch_device)
return {"sample": image}
@property
def input_shape(self):
return (3, 9, 16, 16)
@property
def output_shape(self):
return (3, 9, 16, 16)
def prepare_init_args_and_inputs_for_common(self):
init_dict = self.get_autoencoder_kl_magvit_config()
inputs_dict = self.dummy_input
return init_dict, inputs_dict
def test_gradient_checkpointing_is_applied(self):
expected_set = {"EasyAnimateEncoder", "EasyAnimateDecoder"}
super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
@unittest.skip("Not quite sure why this test fails. Revisit later.")
def test_effective_gradient_checkpointing(self):
pass
@unittest.skip("Unsupported test.")
def test_forward_with_norm_groups(self):
pass
@unittest.skip(
"Unsupported test. Error: RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 9 but got size 12 for tensor number 1 in the list."
)
def test_enable_disable_slicing(self):
pass

diffusers/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py

Lines 32 to 70 in 0f1abc4

class AutoencoderKLTemporalDecoderTests(ModelTesterMixin, AutoencoderTesterMixin, unittest.TestCase):
model_class = AutoencoderKLTemporalDecoder
main_input_name = "sample"
base_precision = 1e-2
@property
def dummy_input(self):
batch_size = 3
num_channels = 3
sizes = (32, 32)
image = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
num_frames = 3
return {"sample": image, "num_frames": num_frames}
@property
def input_shape(self):
return (3, 32, 32)
@property
def output_shape(self):
return (3, 32, 32)
def prepare_init_args_and_inputs_for_common(self):
init_dict = {
"block_out_channels": [32, 64],
"in_channels": 3,
"out_channels": 3,
"down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"],
"latent_channels": 4,
"layers_per_block": 2,
}
inputs_dict = self.dummy_input
return init_dict, inputs_dict
def test_gradient_checkpointing_is_applied(self):
expected_set = {"Encoder", "TemporalDecoder", "UNetMidBlock2D"}
super().test_gradient_checkpointing_is_applied(expected_set=expected_set)

diffusers/tests/models/autoencoders/test_models_autoencoder_vidtok.py

Lines 30 to 163 in 0f1abc4

class AutoencoderVidTokTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase):
model_class = AutoencoderVidTok
main_input_name = "sample"
base_precision = 1e-2
def get_autoencoder_vidtok_config(self):
return {
"is_causal": False,
"in_channels": 3,
"out_channels": 3,
"ch": 128,
"ch_mult": [1, 2, 4, 4, 4],
"z_channels": 6,
"double_z": False,
"num_res_blocks": 2,
"regularizer": "fsq",
"codebook_size": 262144,
}
@property
def dummy_input(self):
batch_size = 4
num_frames = 16
num_channels = 3
sizes = (32, 32)
image = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device)
return {"sample": image}
@property
def input_shape(self):
return (3, 16, 32, 32)
@property
def output_shape(self):
return (3, 16, 32, 32)
def prepare_init_args_and_inputs_for_common(self):
init_dict = self.get_autoencoder_vidtok_config()
inputs_dict = self.dummy_input
return init_dict, inputs_dict
def test_enable_disable_tiling(self):
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
torch.manual_seed(0)
model = self.model_class(**init_dict).to(torch_device)
torch.manual_seed(0)
output_without_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0]
torch.manual_seed(0)
model.enable_tiling()
output_with_tiling = model(**inputs_dict, generator=torch.manual_seed(0))[0]
self.assertLess(
(output_without_tiling.detach().cpu().numpy() - output_with_tiling.detach().cpu().numpy()).max(),
0.5,
"VAE tiling should not affect the inference results",
)
torch.manual_seed(0)
model.disable_tiling()
output_without_tiling_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0]
self.assertEqual(
output_without_tiling.detach().cpu().numpy().all(),
output_without_tiling_2.detach().cpu().numpy().all(),
"Without tiling outputs should match with the outputs when tiling is manually disabled.",
)
def test_enable_disable_slicing(self):
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
torch.manual_seed(0)
model = self.model_class(**init_dict).to(torch_device)
inputs_dict.update({"return_dict": False})
torch.manual_seed(0)
output_without_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0]
torch.manual_seed(0)
model.enable_slicing()
output_with_slicing = model(**inputs_dict, generator=torch.manual_seed(0))[0]
self.assertLess(
(output_without_slicing.detach().cpu().numpy() - output_with_slicing.detach().cpu().numpy()).max(),
0.5,
"VAE slicing should not affect the inference results",
)
torch.manual_seed(0)
model.disable_slicing()
output_without_slicing_2 = model(**inputs_dict, generator=torch.manual_seed(0))[0]
self.assertEqual(
output_without_slicing.detach().cpu().numpy().all(),
output_without_slicing_2.detach().cpu().numpy().all(),
"Without slicing outputs should match with the outputs when slicing is manually disabled.",
)
def test_gradient_checkpointing_is_applied(self):
expected_set = {
"VidTokEncoder3D",
"VidTokDecoder3D",
}
super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
def test_forward_with_norm_groups(self):
r"""VidTok uses layernorm instead of groupnorm."""
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
model = self.model_class(**init_dict)
model.to(torch_device)
model.eval()
with torch.no_grad():
output = model(**inputs_dict)
if isinstance(output, dict):
output = output.to_tuple()[0]
self.assertIsNotNone(output)
expected_shape = inputs_dict["sample"].shape
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
@unittest.skip("Unsupported test.")
def test_outputs_equivalence(self):
pass
@unittest.skipIf(IS_GITHUB_ACTIONS, reason="Skipping test inside GitHub Actions environment")
def test_layerwise_casting_training(self):
super().test_layerwise_casting_training()

diffusers/tests/models/autoencoders/test_models_vq.py

Lines 31 to 114 in 0f1abc4

model_class = VQModel
main_input_name = "sample"
@property
def dummy_input(self, sizes=(32, 32)):
batch_size = 4
num_channels = 3
image = floats_tensor((batch_size, num_channels) + sizes).to(torch_device)
return {"sample": image}
@property
def input_shape(self):
return (3, 32, 32)
@property
def output_shape(self):
return (3, 32, 32)
def prepare_init_args_and_inputs_for_common(self):
init_dict = {
"block_out_channels": [8, 16],
"norm_num_groups": 8,
"in_channels": 3,
"out_channels": 3,
"down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"],
"up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"],
"latent_channels": 3,
}
inputs_dict = self.dummy_input
return init_dict, inputs_dict
@unittest.skip("Test not supported.")
def test_forward_signature(self):
pass
@unittest.skip("Test not supported.")
def test_training(self):
pass
def test_from_pretrained_hub(self):
model, loading_info = VQModel.from_pretrained("fusing/vqgan-dummy", output_loading_info=True)
self.assertIsNotNone(model)
self.assertEqual(len(loading_info["missing_keys"]), 0)
model.to(torch_device)
image = model(**self.dummy_input)
assert image is not None, "Make sure output is not None"
def test_output_pretrained(self):
model = VQModel.from_pretrained("fusing/vqgan-dummy")
model.to(torch_device).eval()
torch.manual_seed(0)
backend_manual_seed(torch_device, 0)
image = torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
image = image.to(torch_device)
with torch.no_grad():
output = model(image).sample
output_slice = output[0, -1, -3:, -3:].flatten().cpu()
# fmt: off
expected_output_slice = torch.tensor([-0.0153, -0.4044, -0.1880, -0.5161, -0.2418, -0.4072, -0.1612, -0.0633, -0.0143])
# fmt: on
self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
def test_loss_pretrained(self):
model = VQModel.from_pretrained("fusing/vqgan-dummy")
model.to(torch_device).eval()
torch.manual_seed(0)
backend_manual_seed(torch_device, 0)
image = torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
image = image.to(torch_device)
with torch.no_grad():
output = model(image).commit_loss.cpu()
# fmt: off
expected_output = torch.tensor([0.1936])
# fmt: on
self.assertTrue(torch.allclose(output, expected_output, atol=1e-3))

Problem:
These model test files contain fast coverage but no @slow tests for published checkpoints or expected slices.

Impact:
Checkpoint config/loading/parity regressions can land unnoticed, especially for newer KVAE, VidTok, DC, and Magvit autoencoders.

Reproduction:

from pathlib import Path

checks = {
    "AutoencoderDC": "tests/models/autoencoders/test_models_autoencoder_dc.py",
    "AutoencoderKLKVAE": "tests/models/autoencoders/test_models_autoencoder_kl_kvae.py",
    "AutoencoderKLKVAEVideo": "tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py",
    "AutoencoderKLMagvit": "tests/models/autoencoders/test_models_autoencoder_magvit.py",
    "AutoencoderKLTemporalDecoder": "tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py",
    "AutoencoderVidTok": "tests/models/autoencoders/test_models_autoencoder_vidtok.py",
    "VQModel": "tests/models/autoencoders/test_models_vq.py",
}
for name, file in checks.items():
    print(name, "@slow" in Path(file).read_text())

Relevant precedent:
Existing autoencoder files such as test_models_autoencoder_tiny.py, test_models_autoencoder_oobleck.py, and test_models_consistency_decoder_vae.py include slow integration coverage.

Suggested fix:
Add one slow checkpoint test per missing model where a public or hf-internal-testing checkpoint exists, asserting load, output shape, finite output, and a small expected slice.

Duplicate-search status

Searched GitHub issues and PRs in huggingface/diffusers for model_autoencoders_shared, AutoencoderKLKVAEVideo, KVAEVideo, autoencoder_kl_kvae_video.py, AutoencoderVidTok, AutoencoderTiny, AutoencoderTiny bfloat16, AutoencoderTiny layerwise casting, AutoencoderTiny tiling latent_channels, AutoencoderKL sample_size tuple tiling, and the specific failure modes above. No duplicate issue or PR was found. Broad related matches were not duplicates: open issue #13628 is a Marigold review, PR #11261 added VidTok, and PR #10347 added layerwise casting.

Coverage Status

Public top-level imports, diffusers.models lazy imports, and PyTorch dummy objects are present for the target public classes. Fast tests exist for the target models. Slow model-level test gaps are listed in Issue 7. API docs are present for most target public models; AutoencoderKLTemporalDecoder and AutoencoderVidTok do not appear in the English API model toctree at:

diffusers/docs/source/en/_toctree.yml

Lines 435 to 481 in 0f1abc4

- local: api/models/asymmetricautoencoderkl
title: AsymmetricAutoencoderKL
- local: api/models/autoencoder_dc
title: AutoencoderDC
- local: api/models/autoencoderkl
title: AutoencoderKL
- local: api/models/autoencoderkl_allegro
title: AutoencoderKLAllegro
- local: api/models/autoencoderkl_cogvideox
title: AutoencoderKLCogVideoX
- local: api/models/autoencoderkl_cosmos
title: AutoencoderKLCosmos
- local: api/models/autoencoder_kl_hunyuanimage
title: AutoencoderKLHunyuanImage
- local: api/models/autoencoder_kl_hunyuanimage_refiner
title: AutoencoderKLHunyuanImageRefiner
- local: api/models/autoencoder_kl_hunyuan_video
title: AutoencoderKLHunyuanVideo
- local: api/models/autoencoder_kl_hunyuan_video15
title: AutoencoderKLHunyuanVideo15
- local: api/models/autoencoder_kl_kvae
title: AutoencoderKLKVAE
- local: api/models/autoencoder_kl_kvae_video
title: AutoencoderKLKVAEVideo
- local: api/models/autoencoderkl_audio_ltx_2
title: AutoencoderKLLTX2Audio
- local: api/models/autoencoderkl_ltx_2
title: AutoencoderKLLTX2Video
- local: api/models/autoencoderkl_ltx_video
title: AutoencoderKLLTXVideo
- local: api/models/autoencoderkl_magvit
title: AutoencoderKLMagvit
- local: api/models/autoencoderkl_mochi
title: AutoencoderKLMochi
- local: api/models/autoencoderkl_qwenimage
title: AutoencoderKLQwenImage
- local: api/models/autoencoder_kl_wan
title: AutoencoderKLWan
- local: api/models/autoencoder_rae
title: AutoencoderRAE
- local: api/models/consistency_decoder_vae
title: ConsistencyDecoderVAE
- local: api/models/autoencoder_oobleck
title: Oobleck AutoEncoder
- local: api/models/autoencoder_tiny
title: Tiny AutoEncoder
- local: api/models/vq

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions