remove any text model initialization in weights.py

moondream.py

@@ -75,7 +75,7 @@ class MoondreamModel(nn.Module):
             "vikhyatk/moondream2", revision="2025-01-09"
         )
         self.vision = build_vision_model(config.vision, dtype)
-        self.text = build_text_model(config.text, dtype)
+        self.text = build_text_model(config.text)
 
         # Region Model
         linear_cls = (

weights.py

@@ -1,26 +1,10 @@
 import safetensors
 import torch
 import torch.nn as nn
-import re
 
 from contextlib import contextmanager
 from typing import Callable, List
 
-from .text import build_text_model
-from .config import TextConfig
-
-
-# Our custom linear has an module named linear, so we add linear to the name
-def add_linear_to_key(k: str) -> str:
-    k = k.replace("model.", "")
-    if k.startswith("text.") and ".linear." not in k:
-        k = re.sub(
-            r"(attn\.(?:qkv|proj)|mlp\.fc[12])\.(weight|bias)$",
-            r"\1.linear.\2",
-            k,
-        )
-    return k
-
 
 @contextmanager
 def safetensors_open(safetensors_file: str):
@@ -43,17 +27,12 @@ def safetensors_open(safetensors_file: str):
         yield get_tensor
 
 
-def _load_weights(
-    get_tensor: Callable[[str], torch.Tensor],
-    model: nn.Module,
-    is_quantized: bool = False,
-) -> None:
+def _load_weights(get_tensor: Callable[[str], torch.Tensor], model: nn.Module) -> None:
     """Internal function to load weights using a tensor getter function."""
-    model = model.to(dtype=torch.float16)
+    model = model.to(dtype=torch.bfloat16)
 
     vision = model.vision
     region = model.region
-
     weight_map = {
         "vision_encoder.encoder.model.visual.patch_embed.linear.weight": vision[
             "patch_emb"
@@ -111,42 +90,23 @@ def _load_weights(
             }
         )
 
-    if not is_quantized:
-        for i in range(len(model.text["blocks"])):
-            prefix = f"text_model.transformer.h.{i}"
-            blk = model.text["blocks"][i]
-            weight_map.update(
-                {
-                    f"{prefix}.ln.weight": blk["ln"].weight,
-                    f"{prefix}.ln.bias": blk["ln"].bias,
-                    f"{prefix}.mixer.Wqkv.weight": blk["attn"]["qkv"].weight,
-                    f"{prefix}.mixer.Wqkv.bias": blk["attn"]["qkv"].bias,
-                    f"{prefix}.mixer.out_proj.weight": blk["attn"]["proj"].weight,
-                    f"{prefix}.mixer.out_proj.bias": blk["attn"]["proj"].bias,
-                    f"{prefix}.mlp.fc1.weight": blk["mlp"]["fc1"].weight,
-                    f"{prefix}.mlp.fc1.bias": blk["mlp"]["fc1"].bias,
-                    f"{prefix}.mlp.fc2.weight": blk["mlp"]["fc2"].weight,
-                    f"{prefix}.mlp.fc2.bias": blk["mlp"]["fc2"].bias,
-                }
-            )
-    else:  # add special quantized path. this is specific to how bitblas expects weights to be loaded (.qweight)
-        for i in range(len(model.text["blocks"])):
-            prefix = f"text_model.transformer.h.{i}"
-            blk = model.text["blocks"][i]
-            weight_map.update(
-                {
-                    f"{prefix}.ln.qweight": blk["ln"].weight,
-                    f"{prefix}.ln.bias": blk["ln"].bias,
-                    f"{prefix}.mixer.Wqkv.qweight": blk["attn"]["qkv"].weight,
-                    f"{prefix}.mixer.Wqkv.bias": blk["attn"]["qkv"].bias,
-                    f"{prefix}.mixer.out_proj.qweight": blk["attn"]["proj"].weight,
-                    f"{prefix}.mixer.out_proj.bias": blk["attn"]["proj"].bias,
-                    f"{prefix}.mlp.fc1.qweight": blk["mlp"]["fc1"].weight,
-                    f"{prefix}.mlp.fc1.bias": blk["mlp"]["fc1"].bias,
-                    f"{prefix}.mlp.fc2.qweight": blk["mlp"]["fc2"].weight,
-                    f"{prefix}.mlp.fc2.bias": blk["mlp"]["fc2"].bias,
-                }
-            )
+    for i in range(len(model.text["blocks"])):
+        prefix = f"text_model.transformer.h.{i}"
+        blk = model.text["blocks"][i]
+        weight_map.update(
+            {
+                f"{prefix}.ln.weight": blk["ln"].weight,
+                f"{prefix}.ln.bias": blk["ln"].bias,
+                f"{prefix}.mixer.Wqkv.weight": blk["attn"]["qkv"].weight,
+                f"{prefix}.mixer.Wqkv.bias": blk["attn"]["qkv"].bias,
+                f"{prefix}.mixer.out_proj.weight": blk["attn"]["proj"].weight,
+                f"{prefix}.mixer.out_proj.bias": blk["attn"]["proj"].bias,
+                f"{prefix}.mlp.fc1.weight": blk["mlp"]["fc1"].weight,
+                f"{prefix}.mlp.fc1.bias": blk["mlp"]["fc1"].bias,
+                f"{prefix}.mlp.fc2.weight": blk["mlp"]["fc2"].weight,
+                f"{prefix}.mlp.fc2.bias": blk["mlp"]["fc2"].bias,
+            }
+        )
 
     for key, tensor in weight_map.items():
         tensor.data.copy_(get_tensor(key))
@@ -160,75 +120,37 @@ def _load_weights(
 def load_weights_from_safetensors(weights_file: str, model: nn.Module) -> None:
     """Load weights from a safetensors file into a MoondreamModel instance."""
     with safetensors_open(weights_file) as get_tensor:
-        all_keys = get_tensor.keys()
-
-        is_quantized = any(
-            ".qweight" in key or "_quantized" in key or "quant." in key
-            for key in all_keys
-        )
-
-        if "text_model.transformer.h.0.ln.weight" in all_keys:
-            layernorm_dtype = get_tensor("text_model.transformer.h.0.ln.weight").dtype
-        else:
-            layernorm_dtype = torch.float16
-
-        linear_dtype = torch.int8 if is_quantized else torch.float16
-
-        model.text = build_text_model(
-            TextConfig
-        )
-        if model.setup_caches_flag:
-            model._setup_caches()
-
         if (
-            "vision.blocks.0.attn.proj.bias" in all_keys
-            or "model.vision.blocks.0.attn.proj.bias" in all_keys
+            "vision.blocks.0.attn.proj.bias" in get_tensor.keys()
+            or "model.vision.blocks.0.attn.proj.bias" in get_tensor.keys()
         ):
             with safetensors_open(weights_file) as get_tensor:
-                tensors = {add_linear_to_key(k): get_tensor(k) for k in all_keys}
+                tensors = {
+                    k.replace("model.", ""): get_tensor(k) for k in get_tensor.keys()
+                }
                 model.load_state_dict(tensors, strict=False)
         else:
             # Wrap the get_tensor function to handle key normalization
-            name_map = {k.replace("._orig_mod", ""): k for k in all_keys}
+            name_map = {k.replace("._orig_mod", ""): k for k in get_tensor.keys()}
             _load_weights(
-                lambda x: get_tensor(name_map[x]).to(dtype=torch.float16),
-                model,
-                is_quantized,
+                lambda x: get_tensor(name_map[x]).to(dtype=torch.bfloat16), model
             )
 
 
 def load_weights_from_pt(weights_file: str, model: nn.Module) -> None:
     """Load weights from a PyTorch file into a MoondreamModel instance."""
-    tensors = torch.load(weights_file, map_location="cpu", weights_only=True)
-    all_keys = tensors.keys()
-    is_quantized = any(
-        ".qweight" in key or "_quantized" in key or "quant." in key for key in all_keys
-    )
-
-    if "text.blocks.0.ln.weight" in all_keys:
-        layernorm_dtype = tensors["text.blocks.0.ln.weight"].dtype
-    else:
-        layernorm_dtype = torch.float16
-
-    linear_dtype = torch.int8 if is_quantized else torch.float16
-    model.text = build_text_model(
-        TextConfig
-    )
-    if model.setup_caches_flag:
-        model._setup_caches()
-
-    if (
-        "vision.blocks.0.attn.proj.bias" in all_keys
-        or "model.vision.blocks.0.attn.proj.bias" in all_keys
-    ):
-        tensors = {add_linear_to_key(k): v for k, v in tensors.items()}
-        model.load_state_dict(tensors, strict=False)
+    device = str(torch.empty(0).device)
+    tensors = torch.load(weights_file, map_location=device, weights_only=True)
+    if "vision.blocks.0.attn.proj.bias" in tensors.keys():
+        missing_keys, unexpected_keys = model.load_state_dict(tensors, strict=False)
+        print("Missing keys:", missing_keys)
+        print("Unexpected keys:", unexpected_keys)
     else:
         tensors = {
-            k.replace("._orig_mod", ""): v.to(dtype=torch.float16)
+            k.replace("._orig_mod", ""): v.to(dtype=torch.bfloat16)
             for k, v in tensors.items()
         }
-        _load_weights(lambda x: tensors[x], model, is_quantized)
+        _load_weights(lambda x: tensors[x], model)
 
 
 def load_weights_into_model(weights_file: str, model: nn.Module) -> None: