support pip install

microsoft · Nov 20, 2023 · 4d91608 · 4d91608
1 parent 3d5902f
commit 4d91608
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -0,0 +1,11 @@
+from setuptools import setup, find_namespace_packages
+
+setup(
+    name='vcm',
+    version='0.1',
+    packages=find_namespace_packages(),
+    include_package_data=True,
+    package_data={
+        'vcm': ['video_call_mos_weights.pt'],
+    },    
+)
diff --git a/vcm/__init__.py b/vcm/__init__.py
diff --git a/vcm/model.py b/vcm/model.py
@@ -52,7 +52,13 @@ def _get_frame_freeze_feature(x_ref_frames):
     return x, df_features
 
 class VideoCallMosModel():
-    def __init__(self, checkpoint="vcm/video_call_mos_weights.pt"):
+    def __init__(self, checkpoint=None):
+        if checkpoint is None:
+            try:
+                import pkg_resources
+                checkpoint = pkg_resources.resource_filename('vcm', 'video_call_mos_weights.pt')
+            except pkg_resources.DistributionNotFound:
+                checkpoint = 'vcm/video_call_mos_weights.pt'
         self.dev = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
         print(f"Running on {self.dev}")
         self.model = VcmNet()