Repository: huggingface/nanoVLM
Pull request #17: Safetensor integration with out Mixins
Author: login: ariG23498
Labels: name: enhancement
description: New feature or request
Created: 2025-05-07T12:26:13Z | Merged: 2025-05-08T09:42:12Z

This is in continuation to #15 

CC: @lusxvr @Wauplin


Diff:
diff --git a/.gitignore b/.gitignore
index 35ed71f3..b096b6ff 100644
--- a/.gitignore
+++ b/.gitignore
@@ -19,4 +19,6 @@ logs/
 pyproject.toml
 uv.lock
 
-.DS_Store
\ No newline at end of file
+.DS_Store
+
+checkpoints/
\ No newline at end of file
diff --git a/README.md b/README.md
index e43c9023..728a5e23 100644
--- a/README.md
+++ b/README.md
@@ -11,7 +11,8 @@ nanoVLM is the simplest repository for training/finetuning a small sized Vision-
 Similar to Andrej Karpathy's nanoGPT, we wanted to equip the community with a very simple implementation and training script for Vision Language Models. We do not claim this to be a new SOTA model, rather an educational effort that packs quite a bit of punch if you have the right hardware! You should be able to tweak and play around with the code in no time.
 
 
-### What can nanoVLM do?
+## What can nanoVLM do?
+
 The model definition and training logic of this repository fits in ~750 lines, with some more boilerplate logging and parameter loading. 
 Using the [`SigLIP-B/16-224-85M`](https://huggingface.co/google/siglip-base-patch16-224) and [`HuggingFaceTB/SmolLM2-135M`](https://huggingface.co/HuggingFaceTB/SmolLM2-135M) as backbones results in a **222M** nanoVLM. Training this for ~6h on a single H100 GPU on ~1.7M samples of [the cauldron](https://huggingface.co/datasets/HuggingFaceM4/the_cauldron) results in an accuracy of 35.3% on MMStar.
 
@@ -20,9 +21,12 @@ Using the [`SigLIP-B/16-224-85M`](https://huggingface.co/google/siglip-base-patc
 It is therefore a simple yet powerful platform to get started with VLMs. Perfect to tinker around with different setups and settings, to explore the capabilities and efficiencies of small VLMs!
 
 ## Quick Start
-You can either clone the repository, setup an environment and start with the scripts, or directly [open in Colab](https://colab.research.google.com/github/huggingface/nanoVLM/blob/main/nanoVLM.ipynb). You can also use the [interactive notebook](./nanoVLM.ipynb) to get started!
+
+You can either clone the repository, setup an environement and start with the scripts, or directly [open in Colab](https://colab.research.google.com/github/huggingface/nanoVLM/blob/main/nanoVLM.ipynb). You can also use the [interactive notebook](./nanoVLM.ipynb) to get started!
+
 
 ## Environment Setup
+
 We really like `uv` and recommend using it as your package manager. But feel free to use any one that you prefer.
 
 Let's first clone the repository:
@@ -80,7 +84,59 @@ Generation 4:  This is a cute cat sitting on the surface of the mat. The backgro
 Generation 5:  This is a cat sitting on a rug, which is on the ground. The cat is in brown
 ```
 
+## Hub integration
+
+**nanoVLM** comes with handy methods to load and save the model from the Hugging Face Hub.
+
+### Pretrained weights
+
+Here is how to load from a repo on the Hugging Face Hub. This is the recommended way to start working with the pretrained weights.
+
+```python
+# Load pretrained weights from Hub
+from models.vision_language_model import VisionLanguageModel
+
+model = VisionLanguageModel.from_pretrained("lusxvr/nanoVLM-222M")
+```
+
+### Push to hub
+
+Once you've trained a **nanoVLM** model, you might want to share it on the Hugging Face Hub. You can easily do that with:
+
+```python
+... # Load and train your model
+
+# Push it to `username/my-awesome-nanovlm-model` repo
+model.push_to_hub("my-awesome-nanovlm-model")
+```
+
+The model will be saved on the Hub as a config file `config.json` and a weights file `model.safetensors`. A modelcard `README.md` will also be generated for you with some high-level information. Feel free to update it manually to explain your work.
+
+If the repo does not exist, it will be created for you. By default the repo will be public. You can pass `private=True` if you don't want to share publicly.
+
+
+### Local save/load
+
+If you don't want to host your model on the Hugging Face Hub, it is still possible to save it locally:
+
+```python
+... # Load and train your model
+
+# Save it to a local folder
+model.save_pretrained("path/to/local/model")
+```
+
+You can then reload it from the local path:
+
+```python
+# Load pretrained weights from local path
+from models.vision_language_model import VisionLanguageModel
+
+model = VisionLanguageModel.from_pretrained("path/to/local/model")
+```
+
 ## Citation
+
 If you like the project and want to use it somewhere, please use this citation:
 ```
 @misc{wiedmann2025nanovlm,
diff --git a/generate.py b/generate.py
index 291b66d4..95d56128 100644
--- a/generate.py
+++ b/generate.py
@@ -1,25 +1,21 @@
 import torch
 from PIL import Image
-from huggingface_hub import hf_hub_download
 
 from models.vision_language_model import VisionLanguageModel
-from models.config import VLMConfig
 from data.processors import get_tokenizer, get_image_processor
 
 torch.manual_seed(0)
 
-cfg = VLMConfig()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-# Change to your own model path after training
-path_to_hf_file = hf_hub_download(repo_id="lusxvr/nanoVLM-222M", filename="nanoVLM-222M.pth")
-model = VisionLanguageModel(cfg).to(device)
-model.load_checkpoint(path_to_hf_file)
+# Load model from the Hugging Face Hub or a local directory
+model = VisionLanguageModel.from_pretrained("lusxvr/nanoVLM-222M").to(device)
+
 model.eval()
 
-tokenizer = get_tokenizer(cfg.lm_tokenizer)
-image_processor = get_image_processor(cfg.vit_img_size)
+tokenizer = get_tokenizer(model.cfg.lm_tokenizer)
+image_processor = get_image_processor(model.cfg.vit_img_size)
 
 text = "What is this?"
 template = f"Question: {text} Answer:"
diff --git a/models/config.py b/models/config.py
index 9568873b..c2e52559 100644
--- a/models/config.py
+++ b/models/config.py
@@ -35,12 +35,12 @@ class VLMConfig:
     mp_pixel_shuffle_factor: int = 2
 
     vlm_load_backbone_weights: bool = True
-    vlm_checkpoint_path: str = 'nanoVLM.pth'
+    vlm_checkpoint_path: str = 'checkpoints/nanoVLM-222M'
 
 
 @dataclass
 class TrainConfig:
-    lr_mp: float = 1e-4
+    lr_mp: float = 1e-3
     lr_backbones: float = 5e-5
     data_cutoff_idx: int = None
     batch_size: int = 256
diff --git a/models/language_model.py b/models/language_model.py
index d6397bc6..8c1368f3 100644
--- a/models/language_model.py
+++ b/models/language_model.py
@@ -291,7 +291,7 @@ def from_pretrained(cls, cfg):
         
         # Store original HF vocab size before we modify it
         original_vocab_size = hf_config.vocab_size
-        print(f"Original vocabulary size from pretrained model: {original_vocab_size}")
+        # print(f"Original vocabulary size from pretrained model: {original_vocab_size}")
         
         # Configure model parameters from HF config
         cfg.lm_hidden_dim = hf_config.hidden_size
@@ -303,11 +303,11 @@ def from_pretrained(cls, cfg):
         if hasattr(cfg, 'lm_vocab_size'):
             if cfg.lm_vocab_size < original_vocab_size:
                 raise ValueError(f"Config vocab size ({cfg.lm_vocab_size}) is smaller than pretrained model vocab size ({original_vocab_size})")
-            print(f"Using extended vocabulary size: {cfg.lm_vocab_size}")
+            # print(f"Using vocabulary size: {cfg.lm_vocab_size}")
         else:
             # If not specified, use the original
             cfg.lm_vocab_size = original_vocab_size
-            print(f"Using original vocabulary size: {cfg.lm_vocab_size}")
+            # print(f"Using original vocabulary size: {cfg.lm_vocab_size}")
         
         cfg.lm_n_heads = hf_config.num_attention_heads
         cfg.lm_n_kv_heads = hf_config.num_key_value_heads
@@ -395,7 +395,7 @@ def from_pretrained(cls, cfg):
         # Handle weight tying (if needed)
         if cfg.lm_tie_weights and hasattr(model, 'head') and hasattr(model, 'token_embedding'):
             model.head.weight = model.token_embedding.weight
-            print("Tied token embedding and LM head weights")
+            # print("Tied token embedding and LM head weights")
         
         print(f"Successfully loaded {cfg.lm_model_type} weights from safetensors. Model has {sum(p.numel() for p in model.parameters()):,} parameters.")
         return model
diff --git a/models/vision_language_model.py b/models/vision_language_model.py
index 7a45dee6..87119418 100644
--- a/models/vision_language_model.py
+++ b/models/vision_language_model.py
@@ -1,18 +1,33 @@
+import json
+import os
+import tempfile
+from dataclasses import asdict
+from typing import Optional
+
+
 from models.vision_transformer import ViT
 from models.language_model import LanguageModel
 from models.modality_projector import ModalityProjector
+from models.config import VLMConfig
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from safetensors.torch import load_model, save_model
 
 class VisionLanguageModel(nn.Module):
-    def __init__(self, cfg):
+    def __init__(self, cfg: VLMConfig, load_backbone=True):
         super().__init__()
         self.cfg = cfg
-        self.vision_encoder = ViT(cfg)
-        self.decoder = LanguageModel(cfg)
+        if load_backbone:
+            print("Loading from backbone weights")
+            self.vision_encoder = ViT.from_pretrained(cfg)
+            self.decoder = LanguageModel.from_pretrained(cfg)
+        else:
+            self.vision_encoder = ViT(cfg)
+            self.decoder = LanguageModel(cfg)
         self.MP = ModalityProjector(cfg)
+        self.load_backbone = load_backbone
 
     def forward(self, input_ids, image, attention_mask=None, targets=None):
         image_embd = self.vision_encoder(image)
@@ -93,16 +108,130 @@ def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5):
                 attention_mask = torch.cat((attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)), dim=1)
         
         return generated_tokens
-        
-    def load_checkpoint(self, path):
-        print(f"Loading weights from full VLM checkpoint: {path}")
-        checkpoint = torch.load(path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))        
-        self.load_state_dict(checkpoint)
 
     @classmethod
-    def from_pretrained(cls, cfg):
-        model = cls(cfg)
-        model.vision_encoder = ViT.from_pretrained(cfg)
-        model.decoder = LanguageModel.from_pretrained(cfg)
+    def from_pretrained(
+        cls, repo_id_or_path: str, *, revision: Optional[str] = None
+    ) -> "VisionLanguageModel":
+        """
+        Load a VisionLanguageModel from a local directory or a repo on the Hugging Face Hub.
+
+        Args:
+            repo_id_or_path (str): The path to the local directory or the Hugging Face Hub repo ID.
+
+        Returns:
+            VisionLanguageModel: The loaded model.
+        """
+        # If local folder exists => load from there
+        if os.path.exists(repo_id_or_path):
+            config_path = os.path.join(repo_id_or_path, "config.json")
+            weights_path = os.path.join(repo_id_or_path, "model.safetensors")
+
+            if not os.path.exists(config_path):
+                raise ValueError(
+                    f"Config file not found at {config_path}. Please provide a valid path."
+                )
+            if not os.path.exists(weights_path):
+                raise ValueError(
+                    f"Weights file not found at {weights_path}. Please provide a valid path."
+                )
+        # Otherwise, assume it's a Hugging Face Hub repo
+        else:
+            from huggingface_hub import hf_hub_download
+
+            config_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="config.json", revision=revision
+            )
+            weights_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="model.safetensors", revision=revision
+            )
+
+        # Load config
+        with open(config_path, "r") as f:
+            cfg = VLMConfig(**json.load(f))
+
+        # Initialize model without loading the backbone
+        model = cls(cfg, load_backbone=False)
+
+        # Load safetensors weights
+        load_model(model, weights_path)
+
+        # Done!
+        return model
+
+    def save_pretrained(self, save_directory: str) -> None:
+        """
+        Save the model and configuration to a directory.
+
+        Args:
+            save_directory (str): The directory to save the model and config.
+        """
+        # Create directory if it doesn't exist
+        os.makedirs(save_directory, exist_ok=True)
+
+        # Save config
+        with open(os.path.join(save_directory, "config.json"), "w") as f:
+            f.write(json.dumps(asdict(self.cfg), indent=4))
+
+        # Save weights as safetensors
+        save_model(self, os.path.join(save_directory, "model.safetensors"))
+
+    def push_to_hub(self, repo_id: str, private: bool = False) -> None:
+        """
+        Push the model and configuration to the Hugging Face Hub.
+
+        Args:
+            repo_id (str): The repo ID on the Hugging Face Hub.
+        """
+        from huggingface_hub import create_repo, upload_folder
+
+        # Create repo
+        repo_url = create_repo(repo_id=repo_id, private=private, exist_ok=True)
+        repo_id = repo_url.repo_id
+        print("Created repo: ", repo_url)
+
+        with tempfile.TemporaryDirectory() as save_path:
+            # Save to tmp directory
+            self.save_pretrained(save_path)
+
+            # Save model card
+            with open(os.path.join(save_path, "README.md"), "w") as f:
+                f.write(MODEL_CARD_TEMPLATE.format(repo_id=repo_id))
+
+            # Upload
+            return upload_folder(
+                repo_id=repo_id,
+                repo_type="model",
+                folder_path=save_path,
+                commit_message="Upload nanoVLM using push_to_hub",
+            )
+
+
+MODEL_CARD_TEMPLATE = """
+---
+# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
+# Doc / guide: https://huggingface.co/docs/hub/model-cards
+library_name: nanovlm
+license: mit
+pipeline_tag: image-text-to-text
+tags:
+  - vision-language
+  - multimodal
+  - research
+---
+
+**nanoVLM** is a minimal and lightweight Vision-Language Model (VLM) designed for efficient training and experimentation. Built using pure PyTorch, the entire model architecture and training logic fits within ~750 lines of code. It combines a ViT-based image encoder (SigLIP-B/16-224-85M) with a lightweight causal language model (SmolLM2-135M), resulting in a compact 222M parameter model.
+
+For more information, check out the base model on https://huggingface.co/lusxvr/nanoVLM-222M.
+
+**Usage:**
+
+Clone the nanoVLM repository: https://github.com/huggingface/nanoVLM.
+Follow the install instructions and run the following code:
+
+```python
+from models.vision_language_model import VisionLanguageModel
 
-        return model
\ No newline at end of file
+model = VisionLanguageModel.from_pretrained("{repo_id}")
+```
+"""
diff --git a/models/vision_transformer.py b/models/vision_transformer.py
index 29c4fab1..a9b86c5b 100644
--- a/models/vision_transformer.py
+++ b/models/vision_transformer.py
@@ -188,7 +188,6 @@ def from_pretrained(cls, cfg):
 
         sd = model.state_dict()
         
-
         mapping = {
             'vision_model.embeddings.patch_embedding.weight': 'patch_embedding.conv.weight',
             'vision_model.embeddings.patch_embedding.bias': 'patch_embedding.conv.bias',
diff --git a/nanoVLM.ipynb b/nanoVLM.ipynb
index 0b2e14a7..b3fe58a3 100644
--- a/nanoVLM.ipynb
+++ b/nanoVLM.ipynb
@@ -252,10 +252,7 @@
         "\n",
         "    # Initialize model\n",
         "    if train_cfg.resume_from_vlm_checkpoint:\n",
-        "        model = VisionLanguageModel(vlm_cfg)\n",
-        "        model.load_checkpoint(vlm_cfg.vlm_checkpoint_path)\n",
-        "    elif vlm_cfg.vlm_load_backbone_weights:\n",
-        "        model = VisionLanguageModel.from_pretrained(vlm_cfg)\n",
+        "        model = VisionLanguageModel.from_pretrained(vlm_cfg.vlm_checkpoint_path)\n",
         "    else:\n",
         "        model = VisionLanguageModel(vlm_cfg)\n",
         "\n",
@@ -315,7 +312,7 @@
         "                    epoch_accuracy = test_mmstar(model, tokenizer, test_loader, device)\n",
         "                    if epoch_accuracy > best_accuracy:\n",
         "                      best_accuracy = epoch_accuracy\n",
-        "                      torch.save(getattr(model, '_orig_mod', model).state_dict(), vlm_cfg.vlm_checkpoint_path)\n",
+        "                      model.save_pretrained(save_directory=vlm_cfg.vlm_checkpoint_path)\n",
         "                    print(f\"\\nStep: {global_step}, Loss: {batch_loss:.4f}, Tokens/s: {tokens_per_second:.2f}, Accuracy: {epoch_accuracy:.4f}\")\n",
         "\n",
         "            global_step += 1\n",
@@ -332,7 +329,7 @@
         "\n",
         "    # Summary Statistics\n",
         "    if not train_cfg.eval_in_epochs:\n",
-        "      torch.save(getattr(model, '_orig_mod', model).state_dict(), vlm_cfg.vlm_checkpoint_path)\n",
+        "      model.save_pretrained(save_directory=vlm_cfg.vlm_checkpoint_path)\n",
         "\n",
         "    avg_epoch_time = sum(epoch_times) / len(epoch_times)\n",
         "    total_training_time = sum(epoch_times)\n",
@@ -408,12 +405,12 @@
         "    mp_pixel_shuffle_factor: int = 2\n",
         "\n",
         "    vlm_load_backbone_weights: bool = True\n",
-        "    vlm_checkpoint_path: str = 'nanoVLM.pth'\n",
+        "    vlm_checkpoint_path: str = 'checkpoints/nanoVLM-222M'\n",
         "\n",
         "\n",
         "@dataclass\n",
         "class TrainConfig:\n",
-        "    lr_mp: float = 1e-4\n",
+        "    lr_mp: float = 1e-3\n",
         "    lr_backbones: float = 5e-5\n",
         "    compile: bool = False\n",
         "    data_cutoff_idx: int = 1024 # Let's only use a small subset of the data at first, otherwise it takes very long to see anything :D\n",
diff --git a/train.py b/train.py
index 05b40d53..3356e0e2 100644
--- a/train.py
+++ b/train.py
@@ -125,12 +125,9 @@ def train(train_cfg, vlm_cfg):
 
     # Initialize model
     if train_cfg.resume_from_vlm_checkpoint:
-        model = VisionLanguageModel(vlm_cfg)
-        model.load_checkpoint(vlm_cfg.vlm_checkpoint_path)
-    elif vlm_cfg.vlm_load_backbone_weights:
-        model = VisionLanguageModel.from_pretrained(vlm_cfg)
+        model = VisionLanguageModel.from_pretrained(vlm_cfg.vlm_checkpoint_path)
     else:
-        model = VisionLanguageModel(vlm_cfg)
+        model = VisionLanguageModel(vlm_cfg, load_backbone=vlm_cfg.vlm_load_backbone_weights)
     
     print(f"nanoVLM initialized with {sum(p.numel() for p in model.parameters()):,} parameters") 
     print(f"Training summary: {len(train_loader.dataset)} samples, {len(train_loader)} batches/epoch, batch size {train_cfg.batch_size}")
@@ -186,7 +183,7 @@ def train(train_cfg, vlm_cfg):
                 epoch_accuracy = test_mmstar(model, tokenizer, test_loader, device)
                 if epoch_accuracy > best_accuracy:
                     best_accuracy = epoch_accuracy
-                    torch.save(getattr(model, '_orig_mod', model).state_dict(), vlm_cfg.vlm_checkpoint_path)
+                    model.save_pretrained(save_directory=vlm_cfg.vlm_checkpoint_path)
                     print(f"Step: {global_step}, Loss: {batch_loss:.4f}, Tokens/s: {tokens_per_second:.2f}, Accuracy: {epoch_accuracy:.4f} | Saving checkpoint to {vlm_cfg.vlm_checkpoint_path}")
                 else:
                     print(f"Step: {global_step}, Loss: {batch_loss:.4f}, Tokens/s: {tokens_per_second:.2f}, Accuracy: {epoch_accuracy:.4f}")


Reviews:
user: login: Wauplin
body: (some high level comment following discussion in https://github.com/huggingface/nanoVLM/pull/15)
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T14:14:00Z
user: login: lusxvr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T15:55:59Z
user: login: lusxvr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T16:24:55Z
user: login: Wauplin
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T16:37:38Z
user: login: lusxvr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T17:05:17Z
user: login: Wauplin
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T20:05:19Z
user: login: ariG23498
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T04:50:34Z
user: login: ariG23498
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T05:04:28Z
user: login: lusxvr
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T09:41:06Z
user: login: lusxvr
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T09:41:20Z

Review comments:
diff_hunk: @@ -35,12 +35,12 @@ class VLMConfig:
     mp_pixel_shuffle_factor: int = 2
 
     vlm_load_backbone_weights: bool = True
-    vlm_checkpoint_path: str = 'nanoVLM.pth'
+    vlm_checkpoint_path: str = 'checkpoints/nanoVLM-222M'
 
 
 @dataclass
 class TrainConfig:
-    lr_mp: float = 1e-4
+    lr_mp: float = 1e-3
path: models/config.py
user: login: Wauplin
body: is this change intentional? (seems unrelated to the PR)

```suggestion
    lr_mp: float = 1e-4
```
created_at: 2025-05-07T14:08:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -93,16 +107,131 @@ def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5):
                 attention_mask = torch.cat((attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)), dim=1)
         
         return generated_tokens
-        
-    def load_checkpoint(self, path):
-        print(f"Loading weights from full VLM checkpoint: {path}")
-        checkpoint = torch.load(path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))        
-        self.load_state_dict(checkpoint)
 
     @classmethod
-    def from_pretrained(cls, cfg):
+    def from_pretrained(
+        cls, repo_id_or_path: str, *, revision: Optional[str] = None
+    ) -> "VisionLanguageModel":
+        """
+        Load a VisionLanguageModel from a local directory or a repo on the Hugging Face Hub.
+
+        Args:
+            repo_id_or_path (str): The path to the local directory or the Hugging Face Hub repo ID.
+
+        Returns:
+            VisionLanguageModel: The loaded model.
+        """
+        # If local folder exists => load from there
+        if os.path.exists(repo_id_or_path):
+            config_path = os.path.join(repo_id_or_path, "config.json")
+            weights_path = os.path.join(repo_id_or_path, "model.safetensors")
+
+            if not os.path.exists(config_path):
+                raise ValueError(
+                    f"Config file not found at {config_path}. Please provide a valid path."
+                )
+            if not os.path.exists(weights_path):
+                raise ValueError(
+                    f"Weights file not found at {weights_path}. Please provide a valid path."
+                )
+        # Otherwise, assume it's a Hugging Face Hub repo
+        else:
+            from huggingface_hub import hf_hub_download
+
+            config_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="config.json", revision=revision
+            )
+            weights_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="model.safetensors", revision=revision
+            )
+
+        # Load config
+        with open(config_path, "r") as f:
+            cfg = VLMConfig(**json.load(f))
+
+        # Initialize model
+        cfg.vlm_load_backbone_weights = False
path: models/vision_language_model.py
user: login: Wauplin
body: This config value feels weird to me. From what I understood in the discussions in https://github.com/huggingface/nanoVLM/pull/15#issuecomment-2858131266, whether to load the backbone weights or not should be handled at the method level (`from_pretrained` vs `load_backbone`) and not at the config level. Mutating a config value that have been read from a file can be confusing to the end user.
created_at: 2025-05-07T14:13:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -35,12 +35,12 @@ class VLMConfig:
     mp_pixel_shuffle_factor: int = 2
 
     vlm_load_backbone_weights: bool = True
-    vlm_checkpoint_path: str = 'nanoVLM.pth'
+    vlm_checkpoint_path: str = 'checkpoints/nanoVLM-222M'
 
 
 @dataclass
 class TrainConfig:
-    lr_mp: float = 1e-4
+    lr_mp: float = 1e-3
path: models/config.py
user: login: lusxvr
body: Is an artifact of me testing the training, should not be a big issue
created_at: 2025-05-07T15:55:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -93,16 +107,131 @@ def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5):
                 attention_mask = torch.cat((attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)), dim=1)
         
         return generated_tokens
-        
-    def load_checkpoint(self, path):
-        print(f"Loading weights from full VLM checkpoint: {path}")
-        checkpoint = torch.load(path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))        
-        self.load_state_dict(checkpoint)
 
     @classmethod
-    def from_pretrained(cls, cfg):
+    def from_pretrained(
+        cls, repo_id_or_path: str, *, revision: Optional[str] = None
+    ) -> "VisionLanguageModel":
+        """
+        Load a VisionLanguageModel from a local directory or a repo on the Hugging Face Hub.
+
+        Args:
+            repo_id_or_path (str): The path to the local directory or the Hugging Face Hub repo ID.
+
+        Returns:
+            VisionLanguageModel: The loaded model.
+        """
+        # If local folder exists => load from there
+        if os.path.exists(repo_id_or_path):
+            config_path = os.path.join(repo_id_or_path, "config.json")
+            weights_path = os.path.join(repo_id_or_path, "model.safetensors")
+
+            if not os.path.exists(config_path):
+                raise ValueError(
+                    f"Config file not found at {config_path}. Please provide a valid path."
+                )
+            if not os.path.exists(weights_path):
+                raise ValueError(
+                    f"Weights file not found at {weights_path}. Please provide a valid path."
+                )
+        # Otherwise, assume it's a Hugging Face Hub repo
+        else:
+            from huggingface_hub import hf_hub_download
+
+            config_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="config.json", revision=revision
+            )
+            weights_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="model.safetensors", revision=revision
+            )
+
+        # Load config
+        with open(config_path, "r") as f:
+            cfg = VLMConfig(**json.load(f))
+
+        # Initialize model
+        cfg.vlm_load_backbone_weights = False
path: models/vision_language_model.py
user: login: lusxvr
body: In this case `from_pretrained` is more powerful in the VLM than `load_backbones`. If the user selects / calls `from_pretrained`, it should therefore override any loading of the backbones to make sure there is no interference. But if you think there is a better way to handle this we are happy to adapt
created_at: 2025-05-07T16:24:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -93,16 +107,131 @@ def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5):
                 attention_mask = torch.cat((attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)), dim=1)
         
         return generated_tokens
-        
-    def load_checkpoint(self, path):
-        print(f"Loading weights from full VLM checkpoint: {path}")
-        checkpoint = torch.load(path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))        
-        self.load_state_dict(checkpoint)
 
     @classmethod
-    def from_pretrained(cls, cfg):
+    def from_pretrained(
+        cls, repo_id_or_path: str, *, revision: Optional[str] = None
+    ) -> "VisionLanguageModel":
+        """
+        Load a VisionLanguageModel from a local directory or a repo on the Hugging Face Hub.
+
+        Args:
+            repo_id_or_path (str): The path to the local directory or the Hugging Face Hub repo ID.
+
+        Returns:
+            VisionLanguageModel: The loaded model.
+        """
+        # If local folder exists => load from there
+        if os.path.exists(repo_id_or_path):
+            config_path = os.path.join(repo_id_or_path, "config.json")
+            weights_path = os.path.join(repo_id_or_path, "model.safetensors")
+
+            if not os.path.exists(config_path):
+                raise ValueError(
+                    f"Config file not found at {config_path}. Please provide a valid path."
+                )
+            if not os.path.exists(weights_path):
+                raise ValueError(
+                    f"Weights file not found at {weights_path}. Please provide a valid path."
+                )
+        # Otherwise, assume it's a Hugging Face Hub repo
+        else:
+            from huggingface_hub import hf_hub_download
+
+            config_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="config.json", revision=revision
+            )
+            weights_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="model.safetensors", revision=revision
+            )
+
+        # Load config
+        with open(config_path, "r") as f:
+            cfg = VLMConfig(**json.load(f))
+
+        # Initialize model
+        cfg.vlm_load_backbone_weights = False
path: models/vision_language_model.py
user: login: Wauplin
body: My main point was maybe not clear but let's say someone initialize their models with `vlm_load_backbone_weights=False` (or True), (re)train it and then publish to the Hub, the `vlm_load_backbone_weights` value will be shared as well in the `config.json`. But this config value is not really related to the model, more to how the user wants to use the model _for their use case_. Say a second user wants to initialize from the new repo on the Hub, `vlm_load_backbone_weights` should not have any effect for them.

So my conclusion would be to remove `vlm_load_backbone_weights` altogether and rely only on the methods.
created_at: 2025-05-07T16:37:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -93,16 +107,131 @@ def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5):
                 attention_mask = torch.cat((attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)), dim=1)
         
         return generated_tokens
-        
-    def load_checkpoint(self, path):
-        print(f"Loading weights from full VLM checkpoint: {path}")
-        checkpoint = torch.load(path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))        
-        self.load_state_dict(checkpoint)
 
     @classmethod
-    def from_pretrained(cls, cfg):
+    def from_pretrained(
+        cls, repo_id_or_path: str, *, revision: Optional[str] = None
+    ) -> "VisionLanguageModel":
+        """
+        Load a VisionLanguageModel from a local directory or a repo on the Hugging Face Hub.
+
+        Args:
+            repo_id_or_path (str): The path to the local directory or the Hugging Face Hub repo ID.
+
+        Returns:
+            VisionLanguageModel: The loaded model.
+        """
+        # If local folder exists => load from there
+        if os.path.exists(repo_id_or_path):
+            config_path = os.path.join(repo_id_or_path, "config.json")
+            weights_path = os.path.join(repo_id_or_path, "model.safetensors")
+
+            if not os.path.exists(config_path):
+                raise ValueError(
+                    f"Config file not found at {config_path}. Please provide a valid path."
+                )
+            if not os.path.exists(weights_path):
+                raise ValueError(
+                    f"Weights file not found at {weights_path}. Please provide a valid path."
+                )
+        # Otherwise, assume it's a Hugging Face Hub repo
+        else:
+            from huggingface_hub import hf_hub_download
+
+            config_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="config.json", revision=revision
+            )
+            weights_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="model.safetensors", revision=revision
+            )
+
+        # Load config
+        with open(config_path, "r") as f:
+            cfg = VLMConfig(**json.load(f))
+
+        # Initialize model
+        cfg.vlm_load_backbone_weights = False
path: models/vision_language_model.py
user: login: lusxvr
body: But if we remove the `vlm_load_backbone_weights`, then there would be no way to distinguish a user calling `from_pretrained` to load the backbones for training and a user calling `from_pretrained` to load the full vlm checkpoint for inference. Lets say we remove the parameter and load the backbones by default into the model on initialization. Then for a user that wants to do inference, first the backbone weights get loaded and then directly overwritten by the VLM checkpoint. I think we therefore need this `vlm_load_backbone_weights` but maybe I have thought about this in too many circles now that I am missing something obvious.
created_at: 2025-05-07T17:05:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -93,16 +107,131 @@ def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5):
                 attention_mask = torch.cat((attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)), dim=1)
         
         return generated_tokens
-        
-    def load_checkpoint(self, path):
-        print(f"Loading weights from full VLM checkpoint: {path}")
-        checkpoint = torch.load(path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))        
-        self.load_state_dict(checkpoint)
 
     @classmethod
-    def from_pretrained(cls, cfg):
+    def from_pretrained(
+        cls, repo_id_or_path: str, *, revision: Optional[str] = None
+    ) -> "VisionLanguageModel":
+        """
+        Load a VisionLanguageModel from a local directory or a repo on the Hugging Face Hub.
+
+        Args:
+            repo_id_or_path (str): The path to the local directory or the Hugging Face Hub repo ID.
+
+        Returns:
+            VisionLanguageModel: The loaded model.
+        """
+        # If local folder exists => load from there
+        if os.path.exists(repo_id_or_path):
+            config_path = os.path.join(repo_id_or_path, "config.json")
+            weights_path = os.path.join(repo_id_or_path, "model.safetensors")
+
+            if not os.path.exists(config_path):
+                raise ValueError(
+                    f"Config file not found at {config_path}. Please provide a valid path."
+                )
+            if not os.path.exists(weights_path):
+                raise ValueError(
+                    f"Weights file not found at {weights_path}. Please provide a valid path."
+                )
+        # Otherwise, assume it's a Hugging Face Hub repo
+        else:
+            from huggingface_hub import hf_hub_download
+
+            config_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="config.json", revision=revision
+            )
+            weights_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="model.safetensors", revision=revision
+            )
+
+        # Load config
+        with open(config_path, "r") as f:
+            cfg = VLMConfig(**json.load(f))
+
+        # Initialize model
+        cfg.vlm_load_backbone_weights = False
path: models/vision_language_model.py
user: login: Wauplin
body: That's where a `from_pretrained(..., backbone_only=True)` comes into play => additional parameter set by the user when calling the method, not saved in the model config itself. Allows to have the flexibility you need but not saved to the Hub

(See https://github.com/huggingface/nanoVLM/pull/15#issuecomment-2857963091 for example of what I mean)
created_at: 2025-05-07T20:05:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -93,16 +107,131 @@ def generate(self, input_ids, image, attention_mask=None, max_new_tokens=5):
                 attention_mask = torch.cat((attention_mask, torch.ones((batch_size, 1), device=attention_mask.device)), dim=1)
         
         return generated_tokens
-        
-    def load_checkpoint(self, path):
-        print(f"Loading weights from full VLM checkpoint: {path}")
-        checkpoint = torch.load(path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))        
-        self.load_state_dict(checkpoint)
 
     @classmethod
-    def from_pretrained(cls, cfg):
+    def from_pretrained(
+        cls, repo_id_or_path: str, *, revision: Optional[str] = None
+    ) -> "VisionLanguageModel":
+        """
+        Load a VisionLanguageModel from a local directory or a repo on the Hugging Face Hub.
+
+        Args:
+            repo_id_or_path (str): The path to the local directory or the Hugging Face Hub repo ID.
+
+        Returns:
+            VisionLanguageModel: The loaded model.
+        """
+        # If local folder exists => load from there
+        if os.path.exists(repo_id_or_path):
+            config_path = os.path.join(repo_id_or_path, "config.json")
+            weights_path = os.path.join(repo_id_or_path, "model.safetensors")
+
+            if not os.path.exists(config_path):
+                raise ValueError(
+                    f"Config file not found at {config_path}. Please provide a valid path."
+                )
+            if not os.path.exists(weights_path):
+                raise ValueError(
+                    f"Weights file not found at {weights_path}. Please provide a valid path."
+                )
+        # Otherwise, assume it's a Hugging Face Hub repo
+        else:
+            from huggingface_hub import hf_hub_download
+
+            config_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="config.json", revision=revision
+            )
+            weights_path = hf_hub_download(
+                repo_id=repo_id_or_path, filename="model.safetensors", revision=revision
+            )
+
+        # Load config
+        with open(config_path, "r") as f:
+            cfg = VLMConfig(**json.load(f))
+
+        # Initialize model
+        cfg.vlm_load_backbone_weights = False
path: models/vision_language_model.py
user: login: ariG23498
body: @Wauplin what do you think about `load_backbone` as a parameter in the `__init__` of the VLM class?

This would help the user:
1. Create a randomly initialized model with `load_backbone=False`
2. Create the model and load the backbones with `load_backbone=True`
3. Keep the `from_pretrained` methods to only work to and from the Hub with the entire model in mind (no complexity added)

I think this would make more sense to the user and a mental demarcation would help in the future iterations.
created_at: 2025-05-08T04:50:33Z
author_association: COLLABORATOR
diff_hunk: @@ -1,18 +1,33 @@
+import json
+import os
+import tempfile
+from dataclasses import asdict
+from typing import Optional
+
+
 from models.vision_transformer import ViT
 from models.language_model import LanguageModel
 from models.modality_projector import ModalityProjector
+from models.config import VLMConfig
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from safetensors.torch import load_model, save_model
 
 class VisionLanguageModel(nn.Module):
-    def __init__(self, cfg):
+    def __init__(self, cfg: VLMConfig, load_backbone=True):
path: models/vision_language_model.py
user: login: ariG23498
body: @Wauplin something like this.
created_at: 2025-05-08T05:04:28Z
author_association: COLLABORATOR
diff_hunk: @@ -1,18 +1,33 @@
+import json
+import os
+import tempfile
+from dataclasses import asdict
+from typing import Optional
+
+
 from models.vision_transformer import ViT
 from models.language_model import LanguageModel
 from models.modality_projector import ModalityProjector
+from models.config import VLMConfig
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from safetensors.torch import load_model, save_model
 
 class VisionLanguageModel(nn.Module):
-    def __init__(self, cfg):
+    def __init__(self, cfg: VLMConfig, load_backbone=True):
path: models/vision_language_model.py
user: login: lusxvr
body: LGTM
created_at: 2025-05-08T09:41:20Z
author_association: CONTRIBUTOR

PR comments:
user: login: ariG23498
created_at: 2025-05-07T12:28:17Z
body: I have added `test_<>.py` to check for equivalence. The safetensor bit works.

Once I have the green light, I will quickly move to the following:
1. Generate
2. Train
3. Documentation
author_association: COLLABORATOR
user: login: lusxvr
created_at: 2025-05-07T12:38:05Z
body: LGTM, lets verify the output of the `generate.py` method, this was a problem in the old PR
author_association: CONTRIBUTOR
user: login: ariG23498
created_at: 2025-05-07T12:43:30Z
body: If you run `generate.py` using `ariG23498/nanoVLM-demo` instead of `lusxvr/nanoVLM-222M` the script runs fine.

I am suspecting a wrong conversion of `.pth` weights into safetensors.
author_association: COLLABORATOR
user: login: lusxvr
created_at: 2025-05-07T13:17:32Z
body: Can you remove the `test.py` files?
author_association: CONTRIBUTOR
user: login: ariG23498
created_at: 2025-05-07T13:26:57Z
body: > Can you remove the test.py files?

Do you think it would be beneficial for us to have the test files? This might help us in the later iterations to be correct with equivalence tests.
author_association: COLLABORATOR