| """Strict native-component loader; supply the official decoder source explicitly.""" |
| import hashlib,importlib.util,json,math |
| from pathlib import Path |
| from types import SimpleNamespace |
| from unittest.mock import patch |
| import torch |
| from torch import nn |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
| from diffusers import AutoencoderDC,SanaTransformer2DModel |
| from transformers import Qwen2Config,Qwen2Model,InternVLModel |
| from native.configuration_intern_vit import InternVisionConfig |
| from native.modeling_intern_vit import InternVisionModel |
|
|
| class VisualComponents(nn.Module): |
| def __init__(self,root,decoder_source): |
| super().__init__() |
| root=Path(root) |
| source=Path(decoder_source) |
| expected=json.loads((root/'external_decoder_source.json').read_text()) |
| if hashlib.sha256(source.read_bytes()).hexdigest()!=expected['sha256']: |
| raise ValueError('Decoder source differs from the pinned official source') |
| spec=importlib.util.spec_from_file_location('_unilip_external_vae',source) |
| external=importlib.util.module_from_spec(spec);spec.loader.exec_module(external) |
| cfg=json.loads((root/'source_config.json').read_text()) |
| self.unilip_factor=cfg['unilip_factor'] |
| self.model=nn.Module() |
| m=self.model |
| m.vision_tower=InternVisionModel(InternVisionConfig(**json.loads((root/'intern_vision_config.json').read_text()))) |
| m.multi_modal_projector=nn.Sequential(nn.LayerNorm(4096),nn.Linear(4096,896),nn.GELU(),nn.Linear(896,896)) |
| m.latent_queries=nn.Parameter(torch.empty(1,cfg['n_query'],896)) |
| m.projector=nn.Linear(896,2304) |
| connector_config=Qwen2Config(**cfg['text_config']) |
| connector_config._attn_implementation='eager' |
| m.llm_connector=Qwen2Model(connector_config) |
| del m.llm_connector.layers[:-cfg['connect_layer']] |
| del m.llm_connector.embed_tokens |
| dc_cfg=json.loads((root/'dc_ae_config.json').read_text()) |
| |
| with patch.object(AutoencoderDC,'from_pretrained',side_effect=lambda *a,**k:AutoencoderDC.from_config(dc_cfg)): |
| m.vae_decoder=external.DCAE_Decoder(SimpleNamespace(model=SimpleNamespace(dc_ae_path='configuration-only')),896) |
| m.dit=SanaTransformer2DModel.from_config(json.loads((root/'dit_config.json').read_text())) |
|
|
| def image_features(self,pixels): |
| v=self.model.vision_tower(pixels).last_hidden_state[:,1:] |
| b,l,c=v.shape;s=math.isqrt(l) |
| if s*s!=l or s%2: |
| raise ValueError('Input must yield a square even-sized patch grid') |
| v=v.reshape(b,s,s,c) |
| |
| v=InternVLModel.pixel_shuffle(self,v,scale_factor=0.5) |
| return self.model.multi_modal_projector(v.reshape(b,-1,4*c)) |
|
|
| def encode(self,pixels): |
| return self.model.vae_decoder.clip_down(self.image_features(pixels)) |
|
|
| def decode(self,latent): |
| """Unscaled DC-AE latent. Divide SANA flow latents by unilip_factor first.""" |
| return self.model.vae_decoder.vae_decode(latent) |
|
|
| def condition(self,language_hidden_states,attention_mask=None): |
| b,n,_=language_hidden_states.shape |
| if attention_mask is None: |
| attention_mask=torch.ones(b,n,device=language_hidden_states.device,dtype=torch.bool) |
| full=attention_mask[:,None,:,None].bool() & attention_mask[:,None,None,:].bool() |
| additive=(1-full.to(language_hidden_states.dtype))*(-100000.0) |
| out=self.model.llm_connector(inputs_embeds=language_hidden_states,attention_mask=additive, |
| position_ids=(attention_mask.long().cumsum(-1)-1).clamp_min(0),use_cache=False,return_dict=True) |
| return self.model.projector(out.last_hidden_state) |
|
|
| def load_model(root,decoder_source,device='cpu',dtype=torch.float32): |
| with init_empty_weights(include_buffers=False): |
| m=VisualComponents(root,decoder_source) |
| m.load_state_dict(load_file(str(Path(root)/'model.safetensors')),strict=True,assign=True) |
| return m.to(device=device,dtype=dtype).eval() |
|
|
