component.py

"""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()