| """Strict, component-only loaders. No language model allocation or remote execution.""" |
| import json,pathlib,torch |
| from torch import nn |
| from safetensors.torch import load_file |
| try: |
| from transformers.initialization import no_init_weights |
| except ImportError: |
| from transformers.modeling_utils import no_init_weights |
|
|
| @no_init_weights() |
| def build_model(root,dtype=None,attention='sdpa'): |
| root=pathlib.Path(root);kind=root.name |
| if (root/'pixel_native.py').exists(): |
| from pixel_io import PixelUMMPixelIO |
| model=PixelUMMPixelIO();kind='pixelumm_pixel_io' |
| elif (root/'sensenova').exists(): |
| from pixel_io import SenseNovaPixelIO |
| model=SenseNovaPixelIO(json.loads((root/'source_config.json').read_text()));kind='sensenova_pixel_io' |
| elif (root/'component_config.json').exists(): |
| from mingtok.modeling_mingtok import MingTok,MingTokConfig |
| cfg=json.loads((root/'component_config.json').read_text());model=MingTok(MingTokConfig(**cfg));kind='mingtok' |
| elif (root/'vision_encoder').exists(): |
| from vision_encoder import UniARVisionConfig,UniARVisionModel |
| cfg=UniARVisionConfig.from_dict(json.loads((root/'source_config.json').read_text()));cfg._attn_implementation=attention |
| model=UniARVisionModel(cfg);kind='uniar_bsq' |
| else: |
| import llada_tokenizer as native |
| cfg=json.loads((root/'source_config.json').read_text());vc=native.make_vision_config(cfg);vc._attn_implementation=attention |
| model=nn.Module();model.model=nn.Module();model.model.visual=native.VisionEncoder(vc);model.model.vqmodel=native.VQVAE(native.make_vq_config(cfg));kind='llada_image_tokenizer' |
| if dtype is None:dtype=torch.bfloat16 if kind=='uniar_bsq' else torch.float32 |
| return model.to(dtype=dtype).eval() |
|
|
| def tokenizer_view(model,root): |
| """Expose upstream PIL->token API around strictly loaded LLaDA modules.""" |
| import llada_tokenizer as native |
| t=native.ImageTokenizer.__new__(native.ImageTokenizer) |
| t.device=next(model.parameters()).device;t.dtype=next(model.parameters()).dtype |
| t.visual=model.model.visual;t.vqmodel=model.model.vqmodel |
| t.image_processor=native.ImagePreprocessor(pathlib.Path(root)/'preprocessor_config.json') |
| t.spatial_merge_size=t.visual.spatial_merge_size |
| return t |
|
|
| def load_model(root,dtype=None,attention='sdpa',device='cpu'): |
| root=pathlib.Path(root);model=build_model(root,dtype,attention) |
| state=load_file(str(root/'model.safetensors'),device='cpu') |
| model.load_state_dict(state,strict=True);del state |
| return model.to(device).eval() |
|
|
| @no_init_weights() |
| def load_sigvq(root,dtype=torch.float32,device='cpu'): |
| """Load LLaDA's separate token-to-decoder conditioning bridge with strict state checks.""" |
| from sigvq import SigVQ |
| root=pathlib.Path(root);model=SigVQ(vocab_size=16384,inner_dim=4096).to(dtype=dtype) |
| state=torch.load(root/'sigvq_embedding.pt',map_location='cpu',weights_only=True) |
| model.load_state_dict(state,strict=True);del state |
| return model.to(device).eval() |
|
|
