remove_model_weights_from_wheel

MANIFEST.in

@@ -1,5 +1,6 @@
 include README.md
 include THIRD_PARTY_NOTICES.md
-recursive-include nemotron_page_elements_v3
+recursive-include nemotron_page_elements_v3 *.py *.json *.png
+exclude nemotron_page_elements_v3/weights.pth
 
 

nemotron_page_elements_v3/model.py

@@ -10,8 +10,13 @@ import numpy.typing as npt
 import torch.nn as nn
 import torch.nn.functional as F
 from typing import Dict, List, Tuple, Union
+from huggingface_hub import hf_hub_download
 from nemotron_page_elements_v3.yolox.boxes import postprocess
 
+# HuggingFace repository for downloading model weights
+HF_REPO_ID = "nvidia/nemotron-page-elements-v3"
+WEIGHTS_FILENAME = "nemotron_page_elements_v3/weights.pth"
+
 
 def define_model(config_name: str = "page_element_v3", verbose: bool = True) -> nn.Module:
     """
@@ -32,13 +37,18 @@ def define_model(config_name: str = "page_element_v3", verbose: bool = True) ->
     config = exp_module.Exp()
     model = config.get_model()
 
-    # Load weights
+    # Download weights from HuggingFace Hub (cached locally after first download)
+    if verbose:
+        print(f" -> Downloading/loading weights from HuggingFace: {HF_REPO_ID}")
+
+    weights_path = hf_hub_download(
+        repo_id=HF_REPO_ID,
+        filename=WEIGHTS_FILENAME,
+    )
+    
     if verbose:
-        print(" -> Loading weights from", config.ckpt)
+        print(f" -> Weights cached at: {weights_path}")
 
-    # Find package directory and load weights (nemotron_page_elements_v3)
-    package_dir = os.path.dirname(os.path.abspath(__file__))
-    weights_path = os.path.join(package_dir, "weights.pth")
     state_dict = torch.load(weights_path, map_location="cpu", weights_only=False)
     model.load_state_dict(state_dict["model"], strict=True)
     

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "nemotron-page-elements-v3"
-version = "3.0.0"
+version = "3.0.1"
 description = "NVIDIA Nemotron Page Elements v3: A specialized YOLOX-based object detection model for identifying tables, charts, infographics, titles, headers/footers, and text in document pages"
 readme = "README.md"
 requires-python = ">=3.8"
@@ -58,6 +58,7 @@ dependencies = [
     "tensorboard",
     "pycocotools>=2.0.2",
     "onnx>=1.13.0",
+    "huggingface_hub",
 ]
 
 [project.optional-dependencies]
@@ -83,5 +84,5 @@ where = ["."]
 include = ["nemotron_page_elements_v3*"]
 
 [tool.setuptools.package-data]
-"nemotron_page_elements_v3" = ["*.json", "*.pth", "*.png"]
+"nemotron_page_elements_v3" = ["*.json", "*.png"]