| """Selected active Mobile-O visual branches with strict loading.""" |
| import json |
| from pathlib import Path |
| import torch |
| from torch import nn |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
| from diffusers import AutoencoderDC,SanaTransformer2DModel |
| from native.mobileclip import MCi,load_model_config |
| from native.mobile_block import MobileConditioningProjector |
|
|
| class VisionTower(nn.Module): |
| def __init__(self): |
| super().__init__() |
| cfg=load_model_config('mobileclip_l_1024') |
| self.vision_tower=MCi(cfg['image_cfg']['model_name'],projection_dim=cfg['embed_dim']) |
| def forward(self,pixels): |
| out=self.vision_tower(pixels,return_image_embeddings=True)['image_embeddings'] |
| return out.flatten(2).transpose(1,2) |
|
|
| class VisualComponents(nn.Module): |
| def __init__(self,root): |
| super().__init__() |
| root=Path(root);cfg=json.loads((root/'source_config.json').read_text()) |
| self.model=nn.Module();m=self.model |
| m.vision_tower=VisionTower() |
| m.mm_projector=nn.Sequential(nn.Linear(cfg['mm_hidden_size'],cfg['hidden_size']),nn.GELU(),nn.Linear(cfg['hidden_size'],cfg['hidden_size'])) |
| m.diffusion_connector=MobileConditioningProjector(input_dim=cfg['hidden_size'],hidden_dim=512,output_dim=2304,num_layers=cfg['vlm_num_layers']) |
| m.dit=SanaTransformer2DModel.from_config(json.loads((root/'dit_config.json').read_text())) |
| m.vae=AutoencoderDC.from_config(json.loads((root/'vae_config.json').read_text())) |
|
|
| def image_features(self,pixels): |
| """Pixels after the official resize/pad preprocessing, rescaled to [0,1].""" |
| return self.model.mm_projector(self.model.vision_tower(pixels)) |
|
|
| def condition(self,language_hidden_states): |
| """Tuple/list of original VLM hidden states; final four layers are fused.""" |
| return self.model.diffusion_connector(language_hidden_states) |
|
|
| def encode_latent(self,rgb): |
| return self.model.vae.encode(rgb).latent*self.model.vae.config.scaling_factor |
|
|
| def decode_latent(self,latent): |
| return self.model.vae.decode(latent/self.model.vae.config.scaling_factor).sample |
|
|
| def load_model(root,device='cpu',dtype=torch.float32): |
| with init_empty_weights(include_buffers=False): |
| m=VisualComponents(root) |
| m.load_state_dict(load_file(str(Path(root)/'model.safetensors')),strict=True,assign=True) |
| return m.to(device=device,dtype=dtype).eval() |
|
|
