File size: 1,672 Bytes
21d1ec3 | 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 | """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):
# The original method modifies z in place; preserve caller input here.
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()
|