| """Complete PR-DINOv2 visual tokenizer: dual-input encoder and reconstruction decoder.""" |
| import importlib,json,sys,types |
| from pathlib import Path |
| import torch |
| from torch import nn |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
|
|
| def setup_native(root): |
| path=Path(root).resolve()/'native/stage1' |
| if 'stage1' in sys.modules: |
| if str(path) not in list(sys.modules['stage1'].__path__):raise RuntimeError('Another stage1 implementation is imported; use a fresh process for this pinned package.') |
| else: |
| pkg=types.ModuleType('stage1');pkg.__path__=[str(path)];sys.modules['stage1']=pkg |
|
|
| class VisualTokenizer(nn.Module): |
| def __init__(self,root): |
| super().__init__();root=Path(root).resolve();setup_native(root) |
| from stage1.patch_reparam import PatchReparam |
| c=json.loads((root/'config.json').read_text()) |
| for key in ['encoder_config_path','decoder_config_path','normalization_stat_path']:c[key]=str(root/c[key]) |
| c['encoder_params']['model_name']=str(root/c['encoder_params']['model_name']) |
| self.native=PatchReparam(**c) |
|
|
| def encode(self,rgb): |
| """B,3,H,W RGB float [0,1] -> normalized B,896,16,16 generation grid.""" |
| return self.native.encode(rgb,return_aux=True).merged_tokens_norm |
|
|
| def encode_features(self,rgb): |
| """Return full native semantic/reconstruction feature outputs.""" |
| return self.native.encode(rgb,return_aux=True) |
|
|
| def decode(self,latents,decode_noise=None): |
| """Normalized generation grid -> RGB. None retains original decode_noise=True.""" |
| return self.native.decode(latents,decode_noise=decode_noise) |
|
|
| def forward(self,rgb,decode_noise=None):return self.decode(self.encode(rgb),decode_noise=decode_noise) |
|
|
| def load_model(root,device='cpu'): |
| root=Path(root).resolve() |
| with init_empty_weights(include_buffers=False):wrapper=VisualTokenizer(root) |
| state={} |
| for p in sorted((root/'weights/tokenizer').glob('*.safetensors')): |
| part=load_file(str(p));assert not set(part)&set(state);state.update(part) |
| wrapper.native.load_state_dict(state,strict=True,assign=True);del state |
| return wrapper.to(device=device).eval().requires_grad_(False) |
|
|
