From a96d92ba59c56297445021d7adcd2d3ac12e9ef9 Mon Sep 17 00:00:00 2001 From: Niels Date: Mon, 26 Aug 2024 09:57:03 +0200 Subject: [PATCH] First draft --- NeuFlow/neuflow.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/NeuFlow/neuflow.py b/NeuFlow/neuflow.py index 049f06c..e3e1de5 100644 --- a/NeuFlow/neuflow.py +++ b/NeuFlow/neuflow.py @@ -7,13 +7,14 @@ from NeuFlow import corr from NeuFlow import refine from NeuFlow import upsample -from NeuFlow import utils from NeuFlow import config -import time +from huggingface_hub import PyTorchModelHubMixin -class NeuFlow(torch.nn.Module): +class NeuFlow(torch.nn.Module, + PyTorchModelHubMixin, + repo_url="https://github.com/neufieldrobotics/NeuFlow_v2", license="apache-2.0", pipeline_tag="image-to-image"): def __init__(self): super(NeuFlow, self).__init__()