vikhyatk/moondream2

1import safetensors2import torch3import torch.nn as nn4 5from contextlib import contextmanager6from dataclasses import dataclass7from typing import Callable, List8 9from .layers import AttentionWeights, LayerNormWeights, LinearWeights, MLPWeights10 11 12@dataclass13class VisionBlock:14    ln1: LayerNormWeights15    attn: AttentionWeights16    ln2: LayerNormWeights17    mlp: MLPWeights18 19 20@dataclass21class VisionModel:22    patch_emb: LinearWeights23    pos_emb: torch.Tensor24    blocks: List[VisionBlock]25    post_ln: LayerNormWeights26    proj_mlp: MLPWeights27 28 29@dataclass30class TextBlock:31    ln: LayerNormWeights32    attn: AttentionWeights33    mlp: MLPWeights34 35 36@dataclass37class TextModel:38    wte: torch.Tensor39    blocks: List[TextBlock]40    post_ln: LayerNormWeights41    lm_head: LinearWeights42 43 44@dataclass45class RegionModel:46    coord_features: torch.Tensor47    coord_encoder: LinearWeights48    coord_decoder: MLPWeights49    size_features: torch.Tensor50    size_encoder: LinearWeights51    size_decoder: MLPWeights52 53 54@dataclass55class MoondreamModel:56    vision: VisionModel57    text: TextModel58    region: RegionModel59 60 61@contextmanager62def safetensors_open(safetensors_file: str):63    """64    Simplify interfacing with safetensors files. Eliminates the need to ignore65    type errors when using the `safe_open` function.66    """67    with safetensors.safe_open(68        safetensors_file, framework="pt"69    ) as st:  # pyright: ignore70 71        def get_tensor(name: str) -> torch.Tensor:72            return st.get_tensor(name)73 74        def get_keys() -> List[str]:75            return st.keys()76 77        get_tensor.keys = get_keys78 79        yield get_tensor80 81 82def _load_weights(get_tensor: Callable[[str], torch.Tensor], model: nn.Module) -> None:83    """Internal function to load weights using a tensor getter function."""84    model = model.to(dtype=torch.float16)85 86    # Vision Model87    model.vision["patch_emb"].weight.data.copy_(88        get_tensor("vision_encoder.encoder.model.visual.patch_embed.linear.weight")89    )90    model.vision["patch_emb"].bias.data.copy_(91        get_tensor("vision_encoder.encoder.model.visual.patch_embed.linear.bias")92    )93    model.vision.pos_emb.data.copy_(94        get_tensor("vision_encoder.encoder.model.visual.pos_embed")95    )96 97    for i in range(len(model.vision["blocks"])):98        prefix = f"vision_encoder.encoder.model.visual.blocks.{i}"99 100        # Layer norms101        model.vision["blocks"][i]["ln1"].weight.data.copy_(102            get_tensor(f"{prefix}.norm1.weight")103        )104        model.vision["blocks"][i]["ln1"].bias.data.copy_(105            get_tensor(f"{prefix}.norm1.bias")106        )107        model.vision["blocks"][i]["ln2"].weight.data.copy_(108            get_tensor(f"{prefix}.norm2.weight")109        )110        model.vision["blocks"][i]["ln2"].bias.data.copy_(111            get_tensor(f"{prefix}.norm2.bias")112        )113 114        # Attention115        model.vision["blocks"][i]["attn"]["qkv"].weight.data.copy_(116            get_tensor(f"{prefix}.attn.qkv.weight")117        )118        model.vision["blocks"][i]["attn"]["qkv"].bias.data.copy_(119            get_tensor(f"{prefix}.attn.qkv.bias")120        )121        model.vision["blocks"][i]["attn"]["proj"].weight.data.copy_(122            get_tensor(f"{prefix}.attn.proj.weight")123        )124        model.vision["blocks"][i]["attn"]["proj"].bias.data.copy_(125            get_tensor(f"{prefix}.attn.proj.bias")126        )127 128        # MLP129        model.vision["blocks"][i]["mlp"]["fc1"].weight.data.copy_(130            get_tensor(f"{prefix}.mlp.fc1.weight")131        )132        model.vision["blocks"][i]["mlp"]["fc1"].bias.data.copy_(133            get_tensor(f"{prefix}.mlp.fc1.bias")134        )135        model.vision["blocks"][i]["mlp"]["fc2"].weight.data.copy_(136            get_tensor(f"{prefix}.mlp.fc2.weight")137        )138        model.vision["blocks"][i]["mlp"]["fc2"].bias.data.copy_(139            get_tensor(f"{prefix}.mlp.fc2.bias")140        )141 142    model.vision["post_ln"].weight.data.copy_(143        get_tensor("vision_encoder.encoder.model.visual.norm.weight")144    )145    model.vision["post_ln"].bias.data.copy_(146        get_tensor("vision_encoder.encoder.model.visual.norm.bias")147    )148 149    model.vision["proj_mlp"]["fc1"].weight.data.copy_(150        get_tensor("vision_encoder.projection.mlp.fc1.weight")151    )152    model.vision["proj_mlp"]["fc1"].bias.data.copy_(153        get_tensor("vision_encoder.projection.mlp.fc1.bias")154    )155    model.vision["proj_mlp"]["fc2"].weight.data.copy_(156        get_tensor("vision_encoder.projection.mlp.fc2.weight")157    )158    model.vision["proj_mlp"]["fc2"].bias.data.copy_(159        get_tensor("vision_encoder.projection.mlp.fc2.bias")160    )161 162    # Text Model163    model.text.wte.data.copy_(get_tensor("text_model.transformer.embd.wte.weight"))164 165    for i in range(len(model.text["blocks"])):166        prefix = f"text_model.transformer.h.{i}"167 168        # Layer norm169        model.text["blocks"][i]["ln"].weight.data.copy_(170            get_tensor(f"{prefix}.ln.weight")171        )172        model.text["blocks"][i]["ln"].bias.data.copy_(get_tensor(f"{prefix}.ln.bias"))173 174        # Attention175        model.text["blocks"][i]["attn"]["qkv"].weight.data.copy_(176            get_tensor(f"{prefix}.mixer.Wqkv.weight")177        )178        model.text["blocks"][i]["attn"]["qkv"].bias.data.copy_(179            get_tensor(f"{prefix}.mixer.Wqkv.bias")180        )181        model.text["blocks"][i]["attn"]["proj"].weight.data.copy_(182            get_tensor(f"{prefix}.mixer.out_proj.weight")183        )184        model.text["blocks"][i]["attn"]["proj"].bias.data.copy_(185            get_tensor(f"{prefix}.mixer.out_proj.bias")186        )187 188        # MLP189        model.text["blocks"][i]["mlp"]["fc1"].weight.data.copy_(190            get_tensor(f"{prefix}.mlp.fc1.weight")191        )192        model.text["blocks"][i]["mlp"]["fc1"].bias.data.copy_(193            get_tensor(f"{prefix}.mlp.fc1.bias")194        )195        model.text["blocks"][i]["mlp"]["fc2"].weight.data.copy_(196            get_tensor(f"{prefix}.mlp.fc2.weight")197        )198        model.text["blocks"][i]["mlp"]["fc2"].bias.data.copy_(199            get_tensor(f"{prefix}.mlp.fc2.bias")200        )201 202    model.text["post_ln"].weight.data.copy_(get_tensor("text_model.lm_head.ln.weight"))203    model.text["post_ln"].bias.data.copy_(get_tensor("text_model.lm_head.ln.bias"))204 205    model.text["lm_head"].weight.data.copy_(206        get_tensor("text_model.lm_head.linear.weight")207    )208    model.text["lm_head"].bias.data.copy_(get_tensor("text_model.lm_head.linear.bias"))209 210    # Region Model211    model.region.coord_features.data.copy_(212        get_tensor("region_model.coordinate_features.weight").T213    )214    model.region["coord_encoder"].weight.data.copy_(215        get_tensor("region_model.coordinate_encoder.weight")216    )217    model.region["coord_encoder"].bias.data.copy_(218        get_tensor("region_model.coordinate_encoder.bias")219    )220 221    model.region["coord_decoder"]["fc1"].weight.data.copy_(222        get_tensor("region_model.coordinate_decoder.fc1.weight")223    )224    model.region["coord_decoder"]["fc1"].bias.data.copy_(225        get_tensor("region_model.coordinate_decoder.fc1.bias")226    )227    model.region["coord_decoder"]["fc2"].weight.data.copy_(228        get_tensor("region_model.coordinate_decoder.fc2.weight")229    )230    model.region["coord_decoder"]["fc2"].bias.data.copy_(231        get_tensor("region_model.coordinate_decoder.fc2.bias")232    )233 234    model.region.size_features.data.copy_(235        get_tensor("region_model.size_features.weight").T236    )237    model.region["size_encoder"].weight.data.copy_(238        get_tensor("region_model.size_encoder.weight")239    )240    model.region["size_encoder"].bias.data.copy_(241        get_tensor("region_model.size_encoder.bias")242    )243 244    model.region["size_decoder"]["fc1"].weight.data.copy_(245        get_tensor("region_model.size_decoder.fc1.weight")246    )247    model.region["size_decoder"]["fc1"].bias.data.copy_(248        get_tensor("region_model.size_decoder.fc1.bias")249    )250    model.region["size_decoder"]["fc2"].weight.data.copy_(251        get_tensor("region_model.size_decoder.fc2.weight")252    )253    model.region["size_decoder"]["fc2"].bias.data.copy_(254        get_tensor("region_model.size_decoder.fc2.bias")255    )256 257 258def load_weights_from_safetensors(weights_file: str, model: nn.Module) -> None:259    """Load weights from a safetensors file into a MoondreamModel instance."""260    with safetensors_open(weights_file) as get_tensor:261        # Wrap the get_tensor function to handle key normalization262        name_map = {k.replace("._orig_mod", ""): k for k in get_tensor.keys()}263        _load_weights(lambda x: get_tensor(name_map[x]).to(dtype=torch.float16), model)264 265 266def load_weights_from_pt(weights_file: str, model: nn.Module) -> None:267    """Load weights from a PyTorch file into a MoondreamModel instance."""268    device = str(torch.empty(0).device)269    tensors = torch.load(weights_file, map_location=device, weights_only=True)270    tensors = {271        k.replace("._orig_mod", ""): v.to(dtype=torch.float16)272        for k, v in tensors.items()273    }274    _load_weights(lambda x: tensors[x], model)275 276 277def load_weights_into_model(weights_file: str, model: nn.Module) -> None:278    """279    Load weights from either a safetensors or PyTorch file directly into a MoondreamModel instance.280 281    Args:282        weights_file: Path to weights file (either .safetensors or .pt)283        model: MoondreamModel instance to load weights into284    """285    if weights_file.endswith(".safetensors"):286        load_weights_from_safetensors(weights_file, model)287    else:288        load_weights_from_pt(weights_file, model)289 290    # Make all parameters contiguous291    for param in model.parameters():292        param.data = param.data.contiguous()293