BiliSakura/BitDance-Tokenizer-diffusers

1#!/usr/bin/env python32"""3Test BitDance-Tokenizer-diffusers: load tokenizer autoencoders only (no full inference).4 5Self-contained: uses local bitdance_diffusers (copied from BitDance-14B-64x-diffusers).6"""7import sys8from pathlib import Path9 10import torch11 12# Self-contained: add local path so bitdance_diffusers is found13BASE_DIR = Path(__file__).resolve().parent14sys.path.insert(0, str(BASE_DIR))15 16from bitdance_diffusers import BitDanceAutoencoder17 18REPO = str(BASE_DIR)19SUBFOLDERS = ["ae_d16c32", "ae_d32c128", "ae_d32c256"]20 21print("Loading BitDance-Tokenizer autoencoders...")22for subfolder in SUBFOLDERS:23    ae = BitDanceAutoencoder.from_pretrained(REPO, subfolder=subfolder)24    print(f"  {subfolder}: z_channels={ae.z_channels}, patch_size={ae.patch_size}")25 26# Quick encode/decode test with ae_d16c3227ae = BitDanceAutoencoder.from_pretrained(REPO, subfolder="ae_d16c32")28x = torch.randn(1, 3, 64, 64)29z = ae.encode(x)30y = ae.decode(z)31assert y.shape == x.shape, f"decode shape {y.shape} != input {x.shape}"32print("encode/decode test passed")33