huggingface / huggingface/diffusers

`model_autoencoders_shared` model/pipeline review

オープン
#13,652 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る
主要言語
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

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。