File size: 1,052 Bytes
400716a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | #!/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")
|