huggingface / huggingface/diffusers

`model_autoencoders_shared` model/pipeline review

Open
#13,652 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

# `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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_kl_kvae_video.py#L848-L894

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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_kl_kvae.py#L592-L629

Suggested fix:
```python
# _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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_vidtok.py#L1435-L1488

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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_kl_kvae_video.py#L938-L954

Suggested fix:
```python
if not return_dict:
return (dec,)
return DecoderOutput(sample=dec)
```

## Issue 3: `AutoencoderVidTok` bypasses diffusers attention processors

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_vidtok.py#L426-L447
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_vidtok.py#L938-L979

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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_rae.py#L202-L208
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_rae.py#L393-L442

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_tiny.py#L302-L312
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/autoencoders/test_models_autoencoder_tiny.py#L146-L159

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_tiny.py#L147-L151
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_tiny.py#L190-L196
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_tiny.py#L238-L244

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_kl.py#L132-L162
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_kl_kvae.py#L582-L596
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/consistency_decoder_vae.py#L159-L204

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/autoencoders/test_models_autoencoder_dc.py#L81-L104
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/autoencoders/test_models_autoencoder_kl_kvae.py#L28-L73
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/autoencoders/test_models_autoencoder_kl_kvae_video.py#L28-L118
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/autoencoders/test_models_autoencoder_magvit.py#L28-L97
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/autoencoders/test_models_autoencoder_kl_temporal_decoder.py#L32-L70
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/autoencoders/test_models_autoencoder_vidtok.py#L30-L163
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/autoencoders/test_models_vq.py#L31-L114

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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/docs/source/en/_toctree.yml#L435-L481

Contributor guide

Open the contributing guide

Research direction

Start by separating the seven findings and reading the linked implementations in src/diffusers/models/autoencoders/, especially autoencoder_kl_kvae_video.py, autoencoder_vidtok.py, autoencoder_tiny.py, autoencoder_kl.py, autoencoder_kl_kvae.py, and consistency_decoder_vae.py. Run the supplied reproductions and inspect the corresponding tests under tests/models/autoencoders/. Done means each selected regression is covered by tests and the affected model behavior matches the stated repository conventions.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, testing-qa
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.