| """Pinned BitDance d32c128 binary codec; explicit NCHW [-1,1] interface.""" |
| from pathlib import Path |
| import json |
| import torch |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
| from native.autoencoder import VQModel |
|
|
| def load_model(path=None,device="cpu",dtype=torch.float32): |
| root=Path(path or Path(__file__).parent) |
| with init_empty_weights():model=VQModel(**json.loads((root/"config.json").read_text())) |
| model.load_state_dict(load_file(root/"model.safetensors"),strict=True,assign=True) |
| return model.to(device=device,dtype=dtype).eval() |
|
|
| def pack_bits(latent): |
| """Explicit portable representation: NCHW signs -> NHW x 16 uint8 bytes, little bit order per byte. This is a bundle storage helper, not an upstream AR token-ID mapping.""" |
| if latent.ndim!=4 or latent.shape[1]!=128:raise ValueError("Expected N x 128 x H x W sign latent") |
| if not torch.all((latent==1)|(latent==-1)):raise ValueError("Expected binary signs") |
| bits=(latent>0).permute(0,2,3,1).reshape(*latent.shape[:1],latent.shape[2],latent.shape[3],16,8).to(torch.uint8) |
| return (bits*(2**torch.arange(8,device=bits.device,dtype=torch.int64))).sum(-1).to(torch.uint8) |
|
|
| def unpack_bits(packed,dtype=torch.float32): |
| if packed.dtype!=torch.uint8 or packed.shape[-1]!=16:raise ValueError("Expected NHW x 16 uint8 bytes") |
| bits=((packed.unsqueeze(-1).to(torch.int64)>>torch.arange(8,device=packed.device))&1).reshape(*packed.shape[:-1],128) |
| return (bits.to(dtype)*2-1).permute(0,3,1,2).contiguous() |
|
|
