| """Complete trained non-language Harmon components with native computations.""" |
| 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 einops import rearrange |
| from native.vae import AutoencoderKL |
| from native.mar import mar_base |
|
|
| class HarmonVisualComponents(nn.Module): |
| vae_scale=0.2325 |
| def __init__(self,cfg): |
| super().__init__() |
| self.vae=AutoencoderKL(embed_dim=16,ch_mult=(1,1,2,2,4)) |
| mar=dict(cfg['mar']);assert mar.pop('type')=='mar_base' |
| self.mar=mar_base(**mar) |
| d=cfg['llm']['hidden_size'];s=self.mar.encoder_embed_dim |
| self.proj_in=nn.Sequential(nn.Linear(s,d),nn.SiLU(),nn.Linear(d,d)) |
| self.proj_out=nn.Sequential(nn.Linear(d,d),nn.SiLU(),nn.Linear(d,s)) |
|
|
| def encode(self,x): |
| posterior=self.vae.encode(x) |
| z=posterior.mode().mul_(self.vae_scale) |
| return rearrange(z,'b c (m p) (n q) -> b m n (c p q)',p=self.mar.patch_size,q=self.mar.patch_size) |
|
|
| def decode(self,z): |
| |
| z=z/self.vae_scale |
| z=rearrange(z,'b m n (c p q) -> b c (m p) (n q)',p=self.mar.patch_size,q=self.mar.patch_size) |
| return self.vae.decode(z) |
|
|
| def load_model(path=None,dtype=torch.float32,device='cpu'): |
| root=Path(path or Path(__file__).parent) |
| with init_empty_weights():model=HarmonVisualComponents(json.loads((root/'source_config.json').read_text())) |
| model.load_state_dict(load_file(root/'model.safetensors'),strict=True,assign=True) |
| return model.to(device=device,dtype=dtype).eval() |
|
|
