File size: 2,454 Bytes
d3d5695 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | """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()
|