| #!/usr/bin/env python3 | |
| """ | |
| Test BitDance-Tokenizer-diffusers: load tokenizer autoencoders only (no full inference). | |
| Self-contained: uses local bitdance_diffusers (copied from BitDance-14B-64x-diffusers). | |
| """ | |
| import sys | |
| from pathlib import Path | |
| import torch | |
| # Self-contained: add local path so bitdance_diffusers is found | |
| BASE_DIR = Path(__file__).resolve().parent | |
| sys.path.insert(0, str(BASE_DIR)) | |
| from bitdance_diffusers import BitDanceAutoencoder | |
| REPO = str(BASE_DIR) | |
| SUBFOLDERS = ["ae_d16c32", "ae_d32c128", "ae_d32c256"] | |
| print("Loading BitDance-Tokenizer autoencoders...") | |
| for subfolder in SUBFOLDERS: | |
| ae = BitDanceAutoencoder.from_pretrained(REPO, subfolder=subfolder) | |
| print(f" {subfolder}: z_channels={ae.z_channels}, patch_size={ae.patch_size}") | |
| # Quick encode/decode test with ae_d16c32 | |
| ae = BitDanceAutoencoder.from_pretrained(REPO, subfolder="ae_d16c32") | |
| x = torch.randn(1, 3, 64, 64) | |
| z = ae.encode(x) | |
| y = ae.decode(z) | |
| assert y.shape == x.shape, f"decode shape {y.shape} != input {x.shape}" | |
| print("encode/decode test passed") | |