Skip to content

Commit 54abed9

Browse files
Merge pull request #1 from NielsRogge/add_hf
Add HF integration
2 parents 24bdd7b + a96d92b commit 54abed9

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

NeuFlow/neuflow.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,14 @@
77
from NeuFlow import corr
88
from NeuFlow import refine
99
from NeuFlow import upsample
10-
from NeuFlow import utils
1110
from NeuFlow import config
1211

13-
import time
12+
from huggingface_hub import PyTorchModelHubMixin
1413

1514

16-
class NeuFlow(torch.nn.Module):
15+
class NeuFlow(torch.nn.Module,
16+
PyTorchModelHubMixin,
17+
repo_url="https://github.com/neufieldrobotics/NeuFlow_v2", license="apache-2.0", pipeline_tag="image-to-image"):
1718
def __init__(self):
1819
super(NeuFlow, self).__init__()
1920

0 commit comments

Comments
 (0)