only remove is_quantized

weights.py

@@ -161,6 +161,25 @@ 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
@@ -173,7 +192,8 @@ def load_weights_from_safetensors(weights_file: str, model: nn.Module) -> None:
             name_map = {k.replace("._orig_mod", ""): k for k in all_keys}
             _load_weights(
                 lambda x: get_tensor(name_map[x]).to(dtype=torch.float16),
-                model
+                model,
+                # is_quantized,
             )
 
 
@@ -181,6 +201,21 @@ 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