huggingface / huggingface/diffusers
`model_autoencoders_shared` model/pipeline review
- 主要言語
- Python
- スター
- 34.5k
- フォーク
- 7.3k
- 平均マージ
- 3日 3時間
- マージ済み PR(30日)
- 91
説明
# `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
コントリビューションガイド
調査の方向性
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.
索引モデルが issue の本文から書いたものです。
評価
- 技術スタック
- python, pytorch
- 領域
- machine-learning, testing-qa
- issue の種類
- バグ
- 難易度
- 5/5
- 見積もり時間
- 1週間以上
- 活発さ
- 静か
- 明瞭さ
- おおむね明確
- 初心者へのやさしさ
- 35/100