| """Pinned native visual components. No text backbone or full generation pipeline.""" |
| import json,math |
| from pathlib import Path |
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
| KEY = 'unicom_vision_compression' |
|
|
| def linear_from(state,name): |
| w=state[name+'.weight'];return nn.Linear(w.shape[1],w.shape[0],bias=name+'.bias' in state) |
|
|
| def build(config,s): |
| m=nn.Module() |
| if KEY=='internlumina_u2_visual_io': |
| from native.multicodebook import MultiCodebookARHead |
| m.model=nn.Module();m.model.img_embeddings=nn.ModuleList([nn.Embedding(*s[f'model.img_embeddings.{i}.weight'].shape) for i in range(8)]) |
| m.model.img_embeddings_proj=linear_from(s,'model.img_embeddings_proj') |
| m.lm_img_head=nn.ModuleList([linear_from(s,f'lm_img_head.{i}') for i in range(8)]) |
| m.multi_codebook_head=MultiCodebookARHead(config['hidden_size'],8,channel_depth=6,channel_heads=8) |
| elif KEY=='unicom_vision_compression': |
| from native.siglip2 import Siglip2VisionTransformer,LightProjector |
| from native.compression import AttentionCompressionNet |
| from native.flow_interfaces import TimestepEmbedder,LinearPatchEmbed,LinearFinalLayer |
| m.vision_model=Siglip2VisionTransformer(config['vit']);m.vision_aligner=LightProjector(config['vit_aligner']);m.compression=AttentionCompressionNet(config['siglip_original_dim'],config['siglip_latent_dim']) |
| h=config['hidden_size'];c=config['siglip_latent_dim'];m.timestep_emb=TimestepEmbedder(h);m.time_embed=TimestepEmbedder(h);m.patch_embed=LinearPatchEmbed(c,h);m.final_layer=LinearFinalLayer(h,c) |
| elif KEY=='internvlu_vision_bridge': |
| from native.modeling_intern_vit import InternVisionModel |
| from native.configuration_intern_vit import InternVisionConfig |
| c=InternVisionConfig(**config['vision_config']);c.use_flash_attn=False;m.vision_model=InternVisionModel(c) |
| n=s['mlp1.0.weight'].shape[0];m.mlp1=nn.Sequential(nn.LayerNorm(n),linear_from(s,'mlp1.1'),nn.GELU(),linear_from(s,'mlp1.3')) |
| m.special_token_embedding=nn.Embedding(*s['special_token_embedding.weight'].shape) |
| elif KEY=='hidream_o1_vision_pixel_io': |
| from native.qwen3_vl_transformers import Qwen3VLVisionModel |
| from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLVisionConfig |
| from native.pixel_interfaces import BottleneckPatchEmbed,FinalLayer,TimestepEmbedder |
| c=Qwen3VLVisionConfig(**config['vision_config']);c._attn_implementation='eager';m.model=nn.Module();m.model.visual=Qwen3VLVisionModel(c) |
| h=config['text_config']['hidden_size'];b=s['model.x_embedder.proj1.weight'].shape[0];m.model.x_embedder=BottleneckPatchEmbed(None,32,3,b,h);m.model.t_embedder1=TimestepEmbedder(None,h);m.model.final_layer2=FinalLayer(None,h,32,3) |
| elif KEY=='ming_flash_omni2_vision_io': |
| from native.qwen3_moe_vit import Qwen3MoeVisionTransformer,Qwen3VLMoeVisionConfig |
| c=Qwen3VLMoeVisionConfig(**config['vision_config']);c._attn_implementation='eager';m.vision=Qwen3MoeVisionTransformer(c) |
| m.linear_proj=nn.Sequential(linear_from(s,'linear_proj.0'),nn.GELU(),linear_from(s,'linear_proj.2')) |
| elif KEY=='lance_vision': |
| from native.qwen2_5_vl_vit import Qwen2_5_VisionTransformerPretrainedModel |
| from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig |
| c=Qwen2_5_VLVisionConfig(**config);c._attn_implementation='eager';m=Qwen2_5_VisionTransformerPretrainedModel(c) |
| else:raise ValueError(KEY) |
| return m |
|
|
| def load_model(path=None,dtype=torch.float32,device='cpu'): |
| root=Path(path or Path(__file__).parent);cfg=json.loads((root/'source_config.json').read_text());s=load_file(root/'model.safetensors') |
| with init_empty_weights():m=build(cfg,s) |
| m.load_state_dict(s,strict=True,assign=True) |
| return m.to(device=device,dtype=dtype).eval() |
|
|
| def embed_codes(model,codes): |
| """InternLumina-U2 per-position 8-codebook input. Model vocab also contains mask IDs.""" |
| if KEY!='internlumina_u2_visual_io':raise ValueError(KEY) |
| if codes.shape[-1]!=8:raise ValueError('Expected 8 codebooks in original order') |
| return model.model.img_embeddings_proj(torch.cat([model.model.img_embeddings[i](codes[...,i]) for i in range(8)],dim=-1)) |
|
|
| def codebook_logits(model,hidden,teacher_codes): |
| """Teacher-forced head, without the diffusion backbone. hidden: [positions,2048].""" |
| if KEY!='internlumina_u2_visual_io':raise ValueError(KEY) |
| e=torch.stack([model.model.img_embeddings[i](teacher_codes[:,i]) for i in range(7)],dim=1) |
| x=model.multi_codebook_head(hidden,e) |
| return torch.stack([model.lm_img_head[i](x[:,i]) for i in range(8)],dim=1) |
|
|
| def unicom_features(model,pixel_patches,spatial_shapes,attention_mask=None): |
| """Return 1152-d semantic features and 64-d compressed features, with native NaFlex layout.""" |
| if KEY!='unicom_vision_compression':raise ValueError(KEY) |
| y=model.vision_model(pixel_patches,attention_mask=attention_mask,spatial_shapes=spatial_shapes).last_hidden_state |
| return y,model.compression.compress(y) |
|
|
| def unicom_align(model,features): |
| """Requires the checkpoint's 1216-d aligner input. No speculative feature concatenation.""" |
| if features.shape[-1]!=model.vision_aligner.layers[0].in_features: |
| raise ValueError('Upstream checkpoint aligner expects 1216 features, while raw SigLIP emits 1152. Supply the correct upstream adapter input explicitly.') |
| return model.vision_aligner(features) |
|
|
| def internvlu_features(model,pixels): |
| if KEY!='internvlu_vision_bridge':raise ValueError(KEY) |
| from native.pixel_shuffle import pixel_shuffle_v2 |
| y=model.vision_model(pixel_values=pixels,return_dict=True).last_hidden_state[:,1:] |
| return model.mlp1(pixel_shuffle_v2(None,y,scale_factor=0.5,patch_aspect_ratio=1.0)) |
|
|
| def ming_features(model,patches,grid_thw): |
| if KEY!='ming_flash_omni2_vision_io':raise ValueError(KEY) |
| y=model.vision(patches,grid_thw=grid_thw) |
| if isinstance(y,tuple):y=y[0] |
| return F.normalize(model.linear_proj(y),dim=-1) |
|
|
