From f86775f9b9b22cb810fbe89e898d495f0e8771f5 Mon Sep 17 00:00:00 2001
From: Niels <niels.rogge1@gmail.com>
Date: Mon, 12 May 2025 18:45:00 +0200
Subject: [PATCH] Improve mixin

---
 core/foundation_stereo.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/core/foundation_stereo.py b/core/foundation_stereo.py
index cbbedba..766aefa 100644
--- a/core/foundation_stereo.py
+++ b/core/foundation_stereo.py
@@ -7,6 +7,8 @@
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 
+import argparse
+
 import torch,pdb,logging,timm
 import torch.nn as nn
 import torch.nn.functional as F
@@ -124,8 +126,18 @@ def forward(self, x, features):
 
 
 
-class FoundationStereo(nn.Module, huggingface_hub.PyTorchModelHubMixin):
-    def __init__(self, args):
+class FoundationStereo(nn.Module, huggingface_hub.PyTorchModelHubMixin,
+                       repo_url="https://github.com/NVlabs/FoundationStereo",
+                       paper_url="https://huggingface.co/papers/2501.09898",
+                       pipeline_tag="depth-estimation",
+                       license="other",
+                       coders={
+                            argparse.Namespace : (
+                                lambda x: vars(x),
+                                lambda data: argparse.Namespace(**data),
+                            )
+   }):
+    def __init__(self, args: argparse.Namespace):
         super().__init__()
         self.args = args