remove_weights_from_python_wheel

MANIFEST.in

@@ -1,3 +1,4 @@
 include README.md
 include THIRD_PARTY_NOTICES.md
-recursive-include nemotron_graphic_elements_v1
+recursive-include nemotron_graphic_elements_v1 *.py *.json *.png
+recursive-exclude nemotron_graphic_elements_v1 *.pth

nemotron_graphic_elements_v1/__init__.py

@@ -6,6 +6,8 @@ Nemotron Graphic Elements v1
 
 A specialized object detection system designed to identify and extract key elements 
 from charts and graphs. Based on YOLOX architecture.
+
+Model weights are automatically downloaded from Hugging Face Hub on first use.
 """
 
 __version__ = "1.0.0"

nemotron_graphic_elements_v1/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 .yolox.boxes import postprocess
 
+# HuggingFace repository for downloading model weights
+HF_REPO_ID = "nvidia/nemotron-graphic-elements-v1"
+WEIGHTS_FILENAME = "nemotron_graphic_elements_v1/weights.pth"
+
 
 def define_model(config_name: str = "graphic_element_v1", verbose: bool = True) -> nn.Module:
     """
@@ -30,11 +35,19 @@ def define_model(config_name: str = "graphic_element_v1", verbose: bool = True)
     config = 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}")
 
-    ckpt = torch.load(config.ckpt, map_location="cpu", weights_only=False)
+    ckpt = torch.load(weights_path, map_location="cpu", weights_only=False)
     model.load_state_dict(ckpt["model"], strict=True)
 
     model = YoloXWrapper(model, config)

pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
     "matplotlib>=3.5.0",
     "pandas>=1.3.0",
     "Pillow>=9.0.0",
+    "huggingface_hub>=0.20.0",
 ]
 
 [project.optional-dependencies]
@@ -50,5 +51,5 @@ Repository = "https://huggingface.co/nvidia/nemotron-graphic-elements-v1"
 packages = ["nemotron_graphic_elements_v1", "nemotron_graphic_elements_v1.yolox", "nemotron_graphic_elements_v1.post_processing"]
 
 [tool.setuptools.package-data]
-"nemotron_graphic_elements_v1" = ["*.pth", "*.json", "*.png"]
+"nemotron_graphic_elements_v1" = ["*.json", "*.png"]